Merge pull request #644 from adotkhan/master

Fix OSSH prefix validation

psiphon/common/obfuscator/obfuscator.go

@@ -31,7 +31,6 @@ import (
 	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common"
 	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/errors"
 	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/prng"
-	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/regen"
 	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/transforms"
 	"golang.org/x/crypto/hkdf"
 )
@@ -655,33 +654,11 @@ func makeTerminator(keyword string, prefix []byte, direction string) ([]byte, er
 // with random bytes.
 func makePrefix(spec *OSSHPrefixSpec, keyword, direction string) ([]byte, error) {
 
-	if len(spec.Spec) != 1 || len(spec.Spec[0]) != 2 || spec.Spec[0][1] == "" {
-		return nil, errors.TraceNew("invalid prefix spec")
-	}
-
-	rng := prng.NewPRNGWithSeed(spec.Seed)
-
-	args := &regen.GeneratorArgs{
-		RngSource: rng,
-		ByteMode:  true,
-	}
-
-	gen, err := regen.NewGenerator(spec.Spec[0][1], args)
-	if err != nil {
-		return nil, errors.Trace(err)
-	}
-
-	prefix, err := gen.Generate()
+	prefix, err := spec.Spec.ApplyPrefix(spec.Seed, PREAMBLE_HEADER_LENGTH)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
 
-	if len(prefix) < PREAMBLE_HEADER_LENGTH {
-		// Add random padding to fill up to PREAMBLE_HEADER_LENGTH.
-		padding := rng.Bytes(PREAMBLE_HEADER_LENGTH - len(prefix))
-		prefix = append(prefix, padding...)
-	}
-
 	terminator, err := makeTerminator(keyword, prefix, direction)
 
 	if err != nil {

psiphon/common/parameters/parameters.go

@@ -1117,7 +1117,11 @@ func (p *Parameters) Set(
 					return nil, errors.Trace(err)
 				}
 			case transforms.Specs:
-				err := v.Validate()
+				prefixMode := false
+				if name == OSSHPrefixSpecs || name == ServerOSSHPrefixSpecs {
+					prefixMode = true
+				}
+				err := v.Validate(prefixMode)
 				if err != nil {
 					if skipOnError {
 						continue

psiphon/common/regen/internal_generator.go

@@ -78,7 +78,12 @@ type internalGenerator struct {
 	GenerateFunc func() ([]byte, error)
 }
 
-func (gen *internalGenerator) Generate() ([]byte, error) {
+func (gen *internalGenerator) Generate() (b []byte, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("panicked on bad input: Generate: %v", r)
+		}
+	}()
 	return gen.GenerateFunc()
 }
 

psiphon/common/regen/regen.go

@@ -186,8 +186,8 @@ func (a *GeneratorArgs) initialize() error {
 	}
 
 	if a.MinUnboundedRepeatCount > a.MaxUnboundedRepeatCount {
-		panic(fmt.Sprintf("MinUnboundedRepeatCount(%d) > MaxUnboundedRepeatCount(%d)",
-			a.MinUnboundedRepeatCount, a.MaxUnboundedRepeatCount))
+		return fmt.Errorf("MinUnboundedRepeatCount(%d) > MaxUnboundedRepeatCount(%d)",
+			a.MinUnboundedRepeatCount, a.MaxUnboundedRepeatCount)
 	}
 
 	if a.CaptureGroupHandler == nil {
@@ -199,11 +199,11 @@ func (a *GeneratorArgs) initialize() error {
 
 // Rng returns the random number generator used by generators.
 // Panics if called before the GeneratorArgs has been initialized by NewGenerator.
-func (a *GeneratorArgs) Rng() *rand.Rand {
+func (a *GeneratorArgs) Rng() (*rand.Rand, error) {
 	if a.rng == nil {
-		panic("GeneratorArgs has not been initialized by NewGenerator yet")
+		return nil, fmt.Errorf("GeneratorArgs has not been initialized by NewGenerator yet")
 	}
-	return a.rng
+	return a.rng, nil
 }
 
 // Generator generates random bytes or strings.
@@ -219,7 +219,12 @@ If args is nil, default values are used.
 This function does not seed the default RNG, so you must call rand.Seed() if you want
 non-deterministic strings.
 */
-func GenerateString(pattern string) (string, error) {
+func GenerateString(pattern string) (str string, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("panicked on bad input: GenerateString: %v", r)
+		}
+	}()
 	generator, err := NewGenerator(pattern, nil)
 	if err != nil {
 		return "", err
@@ -237,6 +242,12 @@ func GenerateString(pattern string) (string, error) {
 // character range. This makes it impossible to infer the original negated
 // character class.
 func NewGenerator(pattern string, inputArgs *GeneratorArgs) (generator Generator, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("panicked on bad input: NewGenerator: %v", r)
+		}
+	}()
+
 	args := GeneratorArgs{}
 
 	// Copy inputArgs so the caller can't change them.

psiphon/common/regen/regen_test.go

@@ -188,15 +188,16 @@ func TestGeneratorArgs(t *testing.T) {
 		}
 	})
 
-	t.Run("Panics if repeat bounds are invalid", func(t *testing.T) {
+	t.Run("Error if repeat bounds are invalid", func(t *testing.T) {
 		args := &GeneratorArgs{
 			MinUnboundedRepeatCount: 2,
 			MaxUnboundedRepeatCount: 1,
 		}
 
-		shouldPanicWith(t, func() {
-			_ = args.initialize()
-		}, "MinUnboundedRepeatCount(2) > MaxUnboundedRepeatCount(1)")
+		err := args.initialize()
+		if err.Error() != "MinUnboundedRepeatCount(2) > MaxUnboundedRepeatCount(1)" {
+			t.Fatalf("unexpected error: %v", err)
+		}
 	})
 
 	t.Run("Allow equal repeat bounds", func(t *testing.T) {
@@ -215,12 +216,13 @@ func TestGeneratorArgs(t *testing.T) {
 
 	t.Run("Rng", func(t *testing.T) {
 
-		t.Run("Panics if called before initialize", func(t *testing.T) {
+		t.Run("Error if called before initialize", func(t *testing.T) {
 			args := &GeneratorArgs{}
 
-			shouldPanic(t, func() {
-				_ = args.Rng()
-			})
+			_, err := args.Rng()
+			if err == nil {
+				t.Fatal("expected error")
+			}
 		})
 
 		t.Run("Non-nil after initialize", func(t *testing.T) {
@@ -229,7 +231,10 @@ func TestGeneratorArgs(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
-			rng := args.Rng()
+			rng, err := args.Rng()
+			if err != nil {
+				t.Fatal(err)
+			}
 			if rng == nil {
 				t.Fatal("expected non-nil")
 			}
@@ -906,24 +911,6 @@ func max(values ...int) int {
 	return m
 }
 
-func shouldPanic(t *testing.T, f func()) {
-	t.Helper()
-	defer func() { _ = recover() }()
-	f()
-	t.Errorf("should have panicked")
-}
-
-func shouldPanicWith(t *testing.T, f func(), expected string) {
-	t.Helper()
-	defer func() {
-		if r := recover(); r != expected {
-			t.Errorf("expected panic %q, got %q", expected, r)
-		}
-	}()
-	f()
-	t.Errorf("should have panicked")
-}
-
 func shouldNotPanic(t *testing.T, f func()) {
 	t.Helper()
 	defer func() {

psiphon/common/transforms/transforms.go

@@ -54,16 +54,29 @@ type Specs map[string]Spec
 
 // Validate checks that all entries in a set of Specs is well-formed, with
 // valid regular expressions.
-func (specs Specs) Validate() error {
+func (specs Specs) Validate(prefixMode bool) error {
 	seed, err := prng.NewSeed()
 	if err != nil {
 		return errors.Trace(err)
 	}
+
 	for _, spec := range specs {
+
 		// Call Apply to compile/validate the regular expressions and generators.
-		_, err := spec.ApplyString(seed, "")
-		if err != nil {
-			return errors.Trace(err)
+
+		if prefixMode {
+			if len(spec) != 1 || len(spec[0]) != 2 {
+				return errors.TraceNew("prefix mode requires exactly one transform")
+			}
+			_, err := spec.ApplyPrefix(seed, 0)
+			if err != nil {
+				return errors.Trace(err)
+			}
+		} else {
+			_, err := spec.ApplyString(seed, "")
+			if err != nil {
+				return errors.Trace(err)
+			}
 		}
 	}
 
@@ -142,6 +155,43 @@ func (specs Specs) Select(scope string, scopedSpecs ScopedSpecNames) (string, Sp
 	return specName, spec
 }
 
+// ApplyPrefix unlike other Apply methods, does not apply the Spec to an input.
+// It instead generates a sequence of bytes according to the Spec, and returns
+// at least minLength bytes if the Spec generates fewer than minLength bytes.
+//
+// The input seed is used for all random number generation. The same seed can be
+// supplied to produce the same output, for replay.
+func (spec Spec) ApplyPrefix(seed *prng.Seed, minLength int) ([]byte, error) {
+
+	if len(spec) != 1 || len(spec[0]) != 2 {
+		return nil, errors.TraceNew("prefix mode requires exactly one transform")
+	}
+
+	rng := prng.NewPRNGWithSeed(seed)
+
+	args := &regen.GeneratorArgs{
+		RngSource: rng,
+		ByteMode:  true,
+	}
+	gen, err := regen.NewGenerator(spec[0][1], args)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	prefix, err := gen.Generate()
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	if len(prefix) < minLength {
+		// Add random padding to fill up to minLength.
+		padding := rng.Bytes(minLength - len(prefix))
+		prefix = append(prefix, padding...)
+	}
+
+	return prefix, nil
+}
+
 // ApplyString applies the Spec to the input string, producing the output string.
 //
 // The input seed is used for all random generation. The same seed can be