mirror of https://gitlab.com/yawning/obfs4
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
246 lines
6.7 KiB
Go
246 lines
6.7 KiB
Go
/*
|
|
* Copyright (c) 2014, Yawning Angel <yawning at schwanenlied dot me>
|
|
* All rights reserved.
|
|
*
|
|
* Redistribution and use in source and binary forms, with or without
|
|
* modification, are permitted provided that the following conditions are met:
|
|
*
|
|
* * Redistributions of source code must retain the above copyright notice,
|
|
* this list of conditions and the following disclaimer.
|
|
*
|
|
* * Redistributions in binary form must reproduce the above copyright notice,
|
|
* this list of conditions and the following disclaimer in the documentation
|
|
* and/or other materials provided with the distribution.
|
|
*
|
|
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
|
* ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
|
|
* LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
|
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
|
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
|
* INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
|
* CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
|
* ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
|
* POSSIBILITY OF SUCH DAMAGE.
|
|
*/
|
|
|
|
// Package probdist implements a weighted probability distribution suitable for
|
|
// protocol parameterization. To allow for easy reproduction of a given
|
|
// distribution, the drbg package is used as the random number source.
|
|
package probdist // import "gitlab.com/yawning/obfs4.git/common/probdist"
|
|
|
|
import (
|
|
"bytes"
|
|
"container/list"
|
|
"fmt"
|
|
"math/rand"
|
|
"sync"
|
|
|
|
"gitlab.com/yawning/obfs4.git/common/csrand"
|
|
"gitlab.com/yawning/obfs4.git/common/drbg"
|
|
)
|
|
|
|
const (
|
|
minValues = 1
|
|
maxValues = 100
|
|
)
|
|
|
|
// WeightedDist is a weighted distribution.
|
|
type WeightedDist struct {
|
|
sync.Mutex
|
|
|
|
minValue int
|
|
maxValue int
|
|
biased bool
|
|
values []int
|
|
weights []float64
|
|
|
|
alias []int
|
|
prob []float64
|
|
}
|
|
|
|
// New creates a weighted distribution of values ranging from min to max
|
|
// based on a HashDrbg initialized with seed. Optionally, bias the weight
|
|
// generation to match the ScrambleSuit non-uniform distribution from
|
|
// obfsproxy.
|
|
func New(seed *drbg.Seed, min, max int, biased bool) (w *WeightedDist) {
|
|
w = &WeightedDist{minValue: min, maxValue: max, biased: biased}
|
|
|
|
if max <= min {
|
|
panic(fmt.Sprintf("wDist.Reset(): min >= max (%d, %d)", min, max))
|
|
}
|
|
|
|
w.Reset(seed)
|
|
|
|
return
|
|
}
|
|
|
|
// genValues creates a slice containing a random number of random values
|
|
// that when scaled by adding minValue will fall into [min, max].
|
|
func (w *WeightedDist) genValues(rng *rand.Rand) {
|
|
nValues := (w.maxValue + 1) - w.minValue
|
|
values := rng.Perm(nValues)
|
|
if nValues < minValues {
|
|
nValues = minValues
|
|
}
|
|
if nValues > maxValues {
|
|
nValues = maxValues
|
|
}
|
|
nValues = rng.Intn(nValues) + 1
|
|
w.values = values[:nValues]
|
|
}
|
|
|
|
// genBiasedWeights generates a non-uniform weight list, similar to the
|
|
// ScrambleSuit prob_dist module.
|
|
func (w *WeightedDist) genBiasedWeights(rng *rand.Rand) {
|
|
w.weights = make([]float64, len(w.values))
|
|
|
|
culmProb := 0.0
|
|
for i := range w.weights {
|
|
p := (1.0 - culmProb) * rng.Float64()
|
|
w.weights[i] = p
|
|
culmProb += p
|
|
}
|
|
}
|
|
|
|
// genUniformWeights generates a uniform weight list.
|
|
func (w *WeightedDist) genUniformWeights(rng *rand.Rand) {
|
|
w.weights = make([]float64, len(w.values))
|
|
for i := range w.weights {
|
|
w.weights[i] = rng.Float64()
|
|
}
|
|
}
|
|
|
|
// genTables calculates the alias and prob tables used for Vose's Alias method.
|
|
// Algorithm taken from http://www.keithschwarz.com/darts-dice-coins/
|
|
func (w *WeightedDist) genTables() {
|
|
n := len(w.weights)
|
|
var sum float64
|
|
for _, weight := range w.weights {
|
|
sum += weight
|
|
}
|
|
|
|
// Create arrays $Alias$ and $Prob$, each of size $n$.
|
|
alias := make([]int, n)
|
|
prob := make([]float64, n)
|
|
|
|
// Create two worklists, $Small$ and $Large$.
|
|
small := list.New()
|
|
large := list.New()
|
|
|
|
scaled := make([]float64, n)
|
|
for i, weight := range w.weights {
|
|
// Multiply each probability by $n$.
|
|
p_i := weight * float64(n) / sum
|
|
scaled[i] = p_i
|
|
|
|
// For each scaled probability $p_i$:
|
|
if scaled[i] < 1.0 {
|
|
// If $p_i < 1$, add $i$ to $Small$.
|
|
small.PushBack(i)
|
|
} else {
|
|
// Otherwise ($p_i \ge 1$), add $i$ to $Large$.
|
|
large.PushBack(i)
|
|
}
|
|
}
|
|
|
|
// While $Small$ and $Large$ are not empty: ($Large$ might be emptied first)
|
|
for small.Len() > 0 && large.Len() > 0 {
|
|
// Remove the first element from $Small$; call it $l$.
|
|
l := small.Remove(small.Front()).(int)
|
|
// Remove the first element from $Large$; call it $g$.
|
|
g := large.Remove(large.Front()).(int)
|
|
|
|
// Set $Prob[l] = p_l$.
|
|
prob[l] = scaled[l]
|
|
// Set $Alias[l] = g$.
|
|
alias[l] = g
|
|
|
|
// Set $p_g := (p_g + p_l) - 1$. (This is a more numerically stable option.)
|
|
scaled[g] = (scaled[g] + scaled[l]) - 1.0
|
|
|
|
if scaled[g] < 1.0 {
|
|
// If $p_g < 1$, add $g$ to $Small$.
|
|
small.PushBack(g)
|
|
} else {
|
|
// Otherwise ($p_g \ge 1$), add $g$ to $Large$.
|
|
large.PushBack(g)
|
|
}
|
|
}
|
|
|
|
// While $Large$ is not empty:
|
|
for large.Len() > 0 {
|
|
// Remove the first element from $Large$; call it $g$.
|
|
g := large.Remove(large.Front()).(int)
|
|
// Set $Prob[g] = 1$.
|
|
prob[g] = 1.0
|
|
}
|
|
|
|
// While $Small$ is not empty: This is only possible due to numerical instability.
|
|
for small.Len() > 0 {
|
|
// Remove the first element from $Small$; call it $l$.
|
|
l := small.Remove(small.Front()).(int)
|
|
// Set $Prob[l] = 1$.
|
|
prob[l] = 1.0
|
|
}
|
|
|
|
w.prob = prob
|
|
w.alias = alias
|
|
}
|
|
|
|
// Reset generates a new distribution with the same min/max based on a new
|
|
// seed.
|
|
func (w *WeightedDist) Reset(seed *drbg.Seed) {
|
|
// Initialize the deterministic random number generator.
|
|
drbg, _ := drbg.NewHashDrbg(seed)
|
|
rng := rand.New(drbg)
|
|
|
|
w.Lock()
|
|
defer w.Unlock()
|
|
|
|
w.genValues(rng)
|
|
if w.biased {
|
|
w.genBiasedWeights(rng)
|
|
} else {
|
|
w.genUniformWeights(rng)
|
|
}
|
|
w.genTables()
|
|
}
|
|
|
|
// Sample generates a random value according to the distribution.
|
|
func (w *WeightedDist) Sample() int {
|
|
var idx int
|
|
|
|
w.Lock()
|
|
defer w.Unlock()
|
|
|
|
// Generate a fair die roll from an $n$-sided die; call the side $i$.
|
|
i := csrand.Intn(len(w.values))
|
|
// Flip a biased coin that comes up heads with probability $Prob[i]$.
|
|
if csrand.Float64() <= w.prob[i] {
|
|
// If the coin comes up "heads," return $i$.
|
|
idx = i
|
|
} else {
|
|
// Otherwise, return $Alias[i]$.
|
|
idx = w.alias[i]
|
|
}
|
|
|
|
return w.minValue + w.values[idx]
|
|
}
|
|
|
|
// String returns a dump of the distribution table.
|
|
func (w *WeightedDist) String() string {
|
|
var buf bytes.Buffer
|
|
|
|
buf.WriteString("[ ")
|
|
for i, v := range w.values {
|
|
p := w.weights[i]
|
|
if p > 0.01 { // Squelch tiny probabilities.
|
|
buf.WriteString(fmt.Sprintf("%d: %f ", v, p))
|
|
}
|
|
}
|
|
buf.WriteString("]")
|
|
return buf.String()
|
|
}
|