Extend main command

cmd/obfsproxy/main.go

@@ -5,26 +5,82 @@ import (
 	"os"
 
 	"github.com/askolesov/obfsproxy/pkg"
+	"github.com/askolesov/obfsproxy/pkg/codec"
 	"github.com/spf13/cobra"
 )
 
 func main() {
-	var listenAddr, targetAddr string
+	var (
+		listenAddr string
+		targetAddr string
+		isServer   bool
+		isClient   bool
+		key        string
+		redundancy int
+	)
 
 	rootCmd := &cobra.Command{
 		Use:   "obfsproxy",
 		Short: "A simple obfuscating proxy",
-		Run: func(cmd *cobra.Command, args []string) {
-			proxy := pkg.NewProxy(listenAddr, targetAddr)
+		RunE: func(cmd *cobra.Command, args []string) error {
+			// Validate flags
+			if isServer && isClient {
+				return fmt.Errorf("cannot specify both server (-s) and client (-c) modes")
+			}
+			if !isServer && !isClient {
+				isClient = true
+			}
+
+			// Calculate seed from key
+			var seed uint64
+			for _, ch := range key {
+				seed += uint64(ch)
+			}
+
+			// Create codecs
+			xorer, err := codec.NewXorer([]byte(key))
+			if err != nil {
+				return fmt.Errorf("failed to create xorer: %w", err)
+			}
+
+			injector, err := codec.NewInjector(seed, redundancy)
+			if err != nil {
+				return fmt.Errorf("failed to create injector: %w", err)
+			}
+
+			chain, err := codec.NewChain([]codec.Codec{
+				codec.NewInverter(),
+				xorer,
+				injector,
+			})
+			if err != nil {
+				return fmt.Errorf("failed to create codec chain: %w", err)
+			}
+
+			// Create proxy with appropriate transformer
+			var proxy *pkg.Proxy
+			if isServer {
+				proxy = pkg.NewProxy(listenAddr, targetAddr, chain.NewDecoder())
+			} else {
+				proxy = pkg.NewProxy(listenAddr, targetAddr, chain.NewEncoder())
+			}
+
 			if err := proxy.Start(); err != nil {
-				fmt.Println("Error:", err)
+				fmt.Printf("Error: %v\n", err)
 				os.Exit(1)
 			}
+
+			return nil
 		},
 	}
 
+	// Add flags
 	rootCmd.Flags().StringVarP(&listenAddr, "listen", "l", "localhost:8080", "Address to listen on")
 	rootCmd.Flags().StringVarP(&targetAddr, "target", "t", "localhost:80", "Address to forward to")
+	rootCmd.Flags().BoolVarP(&isServer, "server", "s", false, "Run in server mode")
+	rootCmd.Flags().BoolVarP(&isClient, "client", "c", false, "Run in client mode")
+	rootCmd.Flags().StringVarP(&key, "key", "k", "", "Encryption key (required)")
+	rootCmd.Flags().IntVarP(&redundancy, "redundancy", "r", 50, "Redundancy level (0-1000)")
 
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(err)

go.mod

@@ -2,12 +2,16 @@ module github.com/askolesov/obfsproxy
 
 go 1.23.2
 
+require (
+	github.com/spf13/cobra v1.8.1
+	github.com/stretchr/testify v1.9.0
+	golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c
+)
+
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/spf13/cobra v1.8.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	github.com/stretchr/testify v1.9.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -12,6 +12,9 @@ github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c h1:7dEasQXItcW1xKJ2+gg5VOiBnqWrJc+rq0DPKyvvdbY=
+golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c/go.mod h1:NQtJDoLvd6faHhE7m4T/1IY708gDefGGjR/iUW8yQQ8=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/codec/chain.go

@@ -0,0 +1,46 @@
+package codec
+
+import "errors"
+
+type Chain struct {
+	Codecs []Codec
+}
+
+func NewChain(codecs []Codec) (*Chain, error) {
+	if len(codecs) == 0 {
+		return nil, errors.New("at least one codec is required")
+	}
+	return &Chain{Codecs: codecs}, nil
+}
+
+func (c *Chain) NewEncoder() Transformer {
+	encoders := make([]Transformer, len(c.Codecs))
+	for i, codec := range c.Codecs {
+		encoders[i] = codec.NewEncoder()
+	}
+
+	return func(data []byte) []byte {
+		result := data
+		// Apply encoders in forward order
+		for _, encoder := range encoders {
+			result = encoder(result)
+		}
+		return result
+	}
+}
+
+func (c *Chain) NewDecoder() Transformer {
+	decoders := make([]Transformer, len(c.Codecs))
+	for i, codec := range c.Codecs {
+		decoders[i] = codec.NewDecoder()
+	}
+
+	return func(data []byte) []byte {
+		result := data
+		// Apply decoders in reverse order
+		for i := len(decoders) - 1; i >= 0; i-- {
+			result = decoders[i](result)
+		}
+		return result
+	}
+}

pkg/codec/chain_test.go

@@ -0,0 +1,81 @@
+package codec
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestChain(t *testing.T) {
+	tests := []struct {
+		name   string
+		codecs []Codec
+		input  []byte
+	}{
+		{
+			name: "single codec",
+			codecs: []Codec{
+				NewInverter(),
+			},
+			input: []byte("hello"),
+		},
+		{
+			name: "multiple codecs",
+			codecs: []Codec{
+				NewInverter(),
+				must(NewXorer([]byte("key"))),
+				must(NewInjector(42, 100)),
+			},
+			input: []byte("test data"),
+		},
+		{
+			name: "empty input",
+			codecs: []Codec{
+				NewInverter(),
+				must(NewXorer([]byte("key"))),
+			},
+			input: []byte{},
+		},
+		{
+			name: "large input",
+			codecs: []Codec{
+				NewInverter(),
+				must(NewXorer([]byte("test-key"))),
+				must(NewInjector(42, 50)),
+			},
+			input: generateRandomBytes(1000, 42),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			chain, err := NewChain(tt.codecs)
+			require.NoError(t, err)
+
+			encoder := chain.NewEncoder()
+			decoder := chain.NewDecoder()
+
+			// Test full array encode/decode
+			encoded := encoder(tt.input)
+			decoded := decoder(encoded)
+			require.Equal(t, tt.input, decoded)
+
+			// Test chunked encode/decode
+			encodedChunks := transformByChunks(encoder, tt.input, 7)
+			decodedChunks := transformByChunks(decoder, encodedChunks, 13)
+			require.Equal(t, tt.input, decodedChunks)
+		})
+	}
+}
+
+func TestChainValidation(t *testing.T) {
+	_, err := NewChain(nil)
+	require.Error(t, err)
+
+	_, err = NewChain([]Codec{})
+	require.Error(t, err)
+
+	chain, err := NewChain([]Codec{NewInverter()})
+	require.NoError(t, err)
+	require.NotNil(t, chain)
+}

pkg/proxy.go

@@ -4,17 +4,22 @@ import (
 	"io"
 	"log"
 	"net"
+
+	"github.com/askolesov/obfsproxy/pkg/codec"
 )
 
 type Proxy struct {
 	ListenAddr string
 	TargetAddr string
+
+	Transformer codec.Transformer
 }
 
-func NewProxy(listenAddr, targetAddr string) *Proxy {
+func NewProxy(listenAddr, targetAddr string, transformer codec.Transformer) *Proxy {
 	return &Proxy{
-		ListenAddr: listenAddr,
-		TargetAddr: targetAddr,
+		ListenAddr:  listenAddr,
+		TargetAddr:  targetAddr,
+		Transformer: transformer,
 	}
 }
 
@@ -62,12 +67,11 @@ func (p *Proxy) proxy(dst, src net.Conn) {
 			return
 		}
 
-		// Invert bytes
-		for i := 0; i < n; i++ {
-			buf[i] = ^buf[i]
-		}
+		buf = buf[:n]
+
+		buf = p.Transformer(buf)
 
-		_, err = dst.Write(buf[:n])
+		_, err = dst.Write(buf)
 		if err != nil {
 			log.Printf("Error writing to connection: %v", err)
 			return