Change the weighted distribution algorithm be uniform.

The old way was biasted towards the earlier values.  Thanks to asn for
pointing this out and suggesting an alternative.

As an additional tweak, do not reuse the drbg seed when calculating the
IAT distribution, but instead run the seed through SHA256 first, for
extra tinfoil goodness.

5 files changed, 58 insertions, 20 deletions

diff --git a/csrand/csrand.go b/csrand/csrand.go
index 721fb1e..a3299aa 100644
--- a/csrand/csrand.go
+++ b/csrand/csrand.go
@@ -68,7 +68,12 @@ func (r csRandSource) Seed(seed int64) {
 	// No-op.
 }
 
-// Float64 returns, as a float 64, a pesudo random number in [0.0,1.0).
+// Int63n returns, as a int64, a pseudo random number in [0, n).
+func Int63n(n int64) int64 {
+	return CsRand.Int63n(n)
+}
+
+// Float64 returns, as a float64, a pesudo random number in [0.0,1.0).
 func Float64() float64 {
 	return CsRand.Float64()
 }
diff --git a/framing/framing.go b/framing/framing.go
index 66eeff5..48d12c3 100644
--- a/framing/framing.go
+++ b/framing/framing.go
@@ -287,7 +287,7 @@ func (decoder *Decoder) Decode(data []byte, frames *bytes.Buffer) (int, error) {
 	}
 	out, ok := secretbox.Open(data[:0], box[:n], &decoder.nextNonce, &decoder.key)
 	if !ok || decoder.nextLengthInvalid {
-		// When a random lenght is used (on length error) the tag should always
+		// When a random length is used (on length error) the tag should always
 		// mismatch, but be paranoid.
 		return 0, ErrTagMismatch
 	}
diff --git a/obfs4.go b/obfs4.go
index e4c22f8..c780e0c 100644
--- a/obfs4.go
+++ b/obfs4.go
@@ -33,6 +33,7 @@ package obfs4
 
 import (
 	"bytes"
+	"crypto/sha256"
 	"fmt"
 	"io"
 	"math/rand"
@@ -568,7 +569,12 @@ func DialObfs4DialFn(dialFn DialFn, network, address, nodeID, publicKey string,
 	c := new(Obfs4Conn)
 	c.lenProbDist = newWDist(seed, 0, framing.MaximumSegmentLength)
 	if iatObfuscation {
-		c.iatProbDist = newWDist(seed, 0, maxIatDelay)
+		iatSeedSrc := sha256.Sum256(seed.Bytes()[:])
+		iatSeed, err := DrbgSeedFromBytes(iatSeedSrc[:])
+		if err != nil {
+			return nil, err
+		}
+		c.iatProbDist = newWDist(iatSeed, 0, maxIatDelay)
 	}
 	c.conn, err = dialFn(network, address)
 	if err != nil {
@@ -596,6 +602,7 @@ type Obfs4Listener struct {
 	nodeID  *ntor.NodeID
 
 	seed           *DrbgSeed
+	iatSeed        *DrbgSeed
 	iatObfuscation bool
 
 	closeDelayBytes int
@@ -631,7 +638,7 @@ func (l *Obfs4Listener) AcceptObfs4() (*Obfs4Conn, error) {
 	cObfs.listener = l
 	cObfs.lenProbDist = newWDist(l.seed, 0, framing.MaximumSegmentLength)
 	if l.iatObfuscation {
-		cObfs.iatProbDist = newWDist(l.seed, 0, maxIatDelay)
+		cObfs.iatProbDist = newWDist(l.iatSeed, 0, maxIatDelay)
 	}
 	if err != nil {
 		c.Close()
@@ -692,6 +699,14 @@ func ListenObfs4(network, laddr, nodeID, privateKey, seed string, iatObfuscation
 	if err != nil {
 		return nil, err
 	}
+	l.iatObfuscation = iatObfuscation
+	if l.iatObfuscation {
+		iatSeedSrc := sha256.Sum256(l.seed.Bytes()[:])
+		l.iatSeed, err = DrbgSeedFromBytes(iatSeedSrc[:])
+		if err != nil {
+			return nil, err
+		}
+	}
 
 	l.filter, err = newReplayFilter()
 	if err != nil {
@@ -701,7 +716,6 @@ func ListenObfs4(network, laddr, nodeID, privateKey, seed string, iatObfuscation
 	rng := rand.New(newHashDrbg(l.seed))
 	l.closeDelayBytes = rng.Intn(maxCloseDelayBytes)
 	l.closeDelay = rng.Intn(maxCloseDelay)
-	l.iatObfuscation = iatObfuscation
 
 	// Start up the listener.
 	l.listener, err = net.Listen(network, laddr)
diff --git a/packet.go b/packet.go
index 78f6697..61ed981 100644
--- a/packet.go
+++ b/packet.go
@@ -28,6 +28,7 @@
 package obfs4
 
 import (
+	"crypto/sha256"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -182,7 +183,12 @@ func (c *Obfs4Conn) consumeFramedPackets(w io.Writer) (n int, err error) {
 				}
 				c.lenProbDist.reset(seed)
 				if c.iatProbDist != nil {
-					c.iatProbDist.reset(seed)
+					iatSeedSrc := sha256.Sum256(seed.Bytes()[:])
+					iatSeed, err := DrbgSeedFromBytes(iatSeedSrc[:])
+					if err != nil {
+						break
+					}
+					c.iatProbDist.reset(iatSeed)
 				}
 			}
 		default:
diff --git a/weighted_dist.go b/weighted_dist.go
index 04b0d2d..55432b2 100644
--- a/weighted_dist.go
+++ b/weighted_dist.go
@@ -39,6 +39,11 @@ import (
 	"github.com/yawning/obfs4/csrand"
 )
 
+const (
+	minBuckets = 1
+	maxBuckets = 100
+)
+
 // DrbgSeedLength is the length of the hashDrbg seed.
 const DrbgSeedLength = 32
 
@@ -133,10 +138,11 @@ func (drbg *hashDrbg) Seed(seed int64) {
 
 // wDist is a weighted distribution.
 type wDist struct {
-	minValue int
-	maxValue int
-	values   []int
-	buckets  []float64
+	minValue    int
+	maxValue    int
+	values      []int
+	buckets     []int64
+	totalWeight int64
 
 	rng *rand.Rand
 }
@@ -160,11 +166,11 @@ func newWDist(seed *DrbgSeed, min, max int) (w *wDist) {
 // sample generates a random value according to the distribution.
 func (w *wDist) sample() int {
 	retIdx := 0
-	totalProb := 0.0
-	prob := csrand.Float64()
-	for i, bucketProb := range w.buckets {
-		totalProb += bucketProb
-		if prob <= totalProb {
+	var totalWeight int64
+	weight := csrand.Int63n(w.totalWeight)
+	for i, bucketWeight := range w.buckets {
+		totalWeight += bucketWeight
+		if weight <= totalWeight {
 			retIdx = i
 			break
 		}
@@ -181,15 +187,22 @@ func (w *wDist) reset(seed *DrbgSeed) {
 
 	nBuckets := (w.maxValue + 1) - w.minValue
 	w.values = w.rng.Perm(nBuckets)
+	if nBuckets < minBuckets {
+		nBuckets = minBuckets
+	}
+	if nBuckets > maxBuckets {
+		nBuckets = maxBuckets
+	}
+	nBuckets = w.rng.Intn(nBuckets) + 1
 
-	w.buckets = make([]float64, nBuckets)
-	var totalProb float64
+	w.totalWeight = 0
+	w.buckets = make([]int64, nBuckets)
 	for i, _ := range w.buckets {
-		prob := w.rng.Float64() * (1.0 - totalProb)
+		prob := w.rng.Int63n(1000)
 		w.buckets[i] = prob
-		totalProb += prob
+		w.totalWeight += prob
 	}
-	w.buckets[len(w.buckets)-1] = 1.0
+	w.buckets[len(w.buckets)-1] = w.totalWeight
 }
 
 /* vim :set ts=4 sw=4 sts=4 noet : */