Refactor MakeObfs and MakeDeobfs

cmd/ck-client/admin.go

@@ -15,7 +15,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"log"
 	"net"
 
 	"github.com/cbeuw/Cloak/internal/client"
@@ -39,7 +38,7 @@ type administrator struct {
 	adminUID  []byte
 }
 
-func adminHandshake(sta *client.State) *administrator {
+func adminHandshake(sta *client.State) (*administrator, error) {
 	fmt.Println("Enter the ip:port of your server")
 	var addr string
 	fmt.Scanln(&addr)
@@ -48,16 +47,14 @@ func adminHandshake(sta *client.State) *administrator {
 	fmt.Scanln(&b64AdminUID)
 	adminUID, err := base64.StdEncoding.DecodeString(b64AdminUID)
 	if err != nil {
-		log.Println(err)
+		return nil, err
-		return nil
 	}
 
 	sta.UID = adminUID
 
 	remoteConn, err := net.Dial("tcp", addr)
 	if err != nil {
-		log.Println(err)
+		return nil, err
-		return nil
 	}
 
 	clientHello := TLS.ComposeInitHandshake(sta)
@@ -68,15 +65,14 @@ func adminHandshake(sta *client.State) *administrator {
 	for c := 0; c < 3; c++ {
 		_, err = util.ReadTLS(remoteConn, discardBuf)
 		if err != nil {
-			log.Printf("Reading discarded message %v: %v\n", c, err)
+			return nil, err
-			return nil
 		}
 	}
 
 	reply := TLS.ComposeReply()
 	_, err = remoteConn.Write(reply)
 	a := &administrator{remoteConn, adminUID}
-	return a
+	return a, nil
 }
 
 func (a *administrator) getCommand() []byte {

cmd/ck-client/ck-client.go

@@ -81,6 +81,27 @@ func makeRemoteConn(sta *client.State) (net.Conn, error) {
 
 }
 
+func adminPrompt(sta *client.State) error {
+	a, err := adminHandshake(sta)
+	if err != nil {
+		return err
+	}
+	buf := make([]byte, 16000)
+	for {
+		req := a.getCommand()
+		a.adminConn.Write(req)
+		n, err := a.adminConn.Read(buf)
+		if err != nil {
+			return err
+		}
+		resp, err := a.checkAndDecrypt(buf[:n])
+		if err != nil {
+			return err
+		}
+		fmt.Println(string(resp))
+	}
+}
+
 func main() {
 	// Should be 127.0.0.1 to listen to ss-local on this machine
 	var localHost string
@@ -91,8 +112,7 @@ func main() {
 	// The proxy port,should be 443
 	var remotePort string
 	var pluginOpts string
-
+	var isAdmin bool
-	var isAdmin *bool
 
 	log.SetFlags(log.LstdFlags | log.Lshortfile)
 
@@ -111,7 +131,7 @@ func main() {
 		flag.StringVar(&remotePort, "p", "443", "remotePort: proxy port, should be 443")
 		flag.StringVar(&pluginOpts, "c", "ckclient.json", "pluginOpts: path to ckclient.json or options seperated with semicolons")
 		askVersion := flag.Bool("v", false, "Print the version number")
-		isAdmin = flag.Bool("a", false, "Admin mode")
+		isAdmin = *flag.Bool("a", false, "Admin mode")
 		printUsage := flag.Bool("h", false, "Print this message")
 		flag.Parse()
 
@@ -140,25 +160,14 @@ func main() {
 		log.Fatal(err)
 	}
 
-	if *isAdmin {
+	if isAdmin {
-		a := adminHandshake(sta)
+		err = adminPrompt(sta)
-		buf := make([]byte, 16000)
+		if err != nil {
-		for {
+			log.Println(err)
-			req := a.getCommand()
-			a.adminConn.Write(req)
-			n, err := a.adminConn.Read(buf)
-			if err != nil {
-				log.Println(err)
-				return
-			}
-			resp, err := a.checkAndDecrypt(buf[:n])
-			if err != nil {
-				log.Println(err)
-			}
-			fmt.Println(string(resp))
 		}
 		return
 	}
+
 	if sta.SS_LOCAL_PORT == "" {
 		log.Fatal("Must specify localPort")
 	}
@@ -171,8 +180,8 @@ func main() {
 
 	var UNLIMITED int64 = 1e12
 	valve := mux.MakeValve(1e12, 1e12, &UNLIMITED, &UNLIMITED)
-	obfs := util.MakeObfs(sta.UID)
+	obfs := mux.MakeObfs(sta.UID)
-	deobfs := util.MakeDeobfs(sta.UID)
+	deobfs := mux.MakeDeobfs(sta.UID)
 	sesh := mux.MakeSession(0, valve, obfs, deobfs, util.ReadTLS)
 
 	var wg sync.WaitGroup
@@ -205,17 +214,22 @@ func main() {
 			data := make([]byte, 10240)
 			i, err := io.ReadAtLeast(ssConn, data, 1)
 			if err != nil {
+				log.Println(err)
 				ssConn.Close()
 				return
 			}
 			stream, err := sesh.OpenStream()
 			if err != nil {
+				log.Println(err)
 				ssConn.Close()
 				return
 			}
 			_, err = stream.Write(data[:i])
 			if err != nil {
 				log.Println(err)
+				ssConn.Close()
+				stream.Close()
+				return
 			}
 			go pipe(ssConn, stream)
 			pipe(stream, ssConn)

cmd/ck-server/ck-server.go

@@ -143,7 +143,7 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 		}
 	}
 
-	if sesh, existing := user.GetOrCreateSession(sessionID, util.MakeObfs(UID), util.MakeDeobfs(UID), util.ReadTLS); existing {
+	if sesh, existing := user.GetOrCreateSession(sessionID, mux.MakeObfs(UID), mux.MakeDeobfs(UID), util.ReadTLS); existing {
 		sesh.AddConnection(conn)
 		return
 	} else {

internal/multiplex/obfs.go

@@ -1,12 +1,15 @@
-package util
+package multiplex
 
 import (
 	"encoding/binary"
+	"errors"
 
 	xxhash "github.com/OneOfOne/xxhash"
-	mux "github.com/cbeuw/Cloak/internal/multiplex"
 )
 
+type Obfser func(*Frame) ([]byte, error)
+type Deobfser func([]byte) (*Frame, error)
+
 // For each frame, the three parts of the header is xored with three keys.
 // The keys are generated from the SID and the payload of the frame.
 // FIXME: this code will panic if len(data)<18.
@@ -23,8 +26,11 @@ func genXorKeys(secret []byte, data []byte) (i uint32, ii uint32, iii uint32) {
 	return ret[0], ret[1], ret[2]
 }
 
-func MakeObfs(key []byte) func(*mux.Frame) []byte {
+func MakeObfs(key []byte) Obfser {
-	obfs := func(f *mux.Frame) []byte {
+	obfs := func(f *Frame) ([]byte, error) {
+		if len(f.Payload) < 18 {
+			return nil, errors.New("Payload cannot be shorter than 18 bytes")
+		}
 		obfsedHeader := make([]byte, 12)
 		// header: [StreamID 4 bytes][Seq 4 bytes][Closing 4 bytes]
 		i, ii, iii := genXorKeys(key, f.Payload[0:18])
@@ -42,13 +48,16 @@ func MakeObfs(key []byte) func(*mux.Frame) []byte {
 		copy(obfsed[5:17], obfsedHeader)
 		copy(obfsed[17:], f.Payload)
 		// obfsed: [record layer 5 bytes][cipherheader 12 bytes][payload]
-		return obfsed
+		return obfsed, nil
 	}
 	return obfs
 }
 
-func MakeDeobfs(key []byte) func([]byte) *mux.Frame {
+func MakeDeobfs(key []byte) Deobfser {
-	deobfs := func(in []byte) *mux.Frame {
+	deobfs := func(in []byte) (*Frame, error) {
+		if len(in) < 30 {
+			return nil, errors.New("Input cannot be shorter than 30 bytes")
+		}
 		peeled := in[5:]
 		i, ii, iii := genXorKeys(key, peeled[12:30])
 		streamID := binary.BigEndian.Uint32(peeled[0:4]) ^ i
@@ -56,13 +65,13 @@ func MakeDeobfs(key []byte) func([]byte) *mux.Frame {
 		closing := binary.BigEndian.Uint32(peeled[8:12]) ^ iii
 		payload := make([]byte, len(peeled)-12)
 		copy(payload, peeled[12:])
-		ret := &mux.Frame{
+		ret := &Frame{
 			StreamID: streamID,
 			Seq:      seq,
 			Closing:  closing,
 			Payload:  payload,
 		}
-		return ret
+		return ret, nil
 	}
 	return deobfs
 }

internal/multiplex/session.go

@@ -20,9 +20,9 @@ type Session struct {
 	id uint32 // This field isn't acutally used
 
 	// Used in Stream.Write. Add multiplexing headers, encrypt and add TLS header
-	obfs func(*Frame) []byte
+	obfs Obfser
 	// Remove TLS header, decrypt and unmarshall multiplexing headers
-	deobfs func([]byte) *Frame
+	deobfs Deobfser
 	// This is supposed to read one TLS message, the same as GoQuiet's ReadTillDrain
 	obfsedRead func(net.Conn, []byte) (int, error)
 
@@ -43,7 +43,7 @@ type Session struct {
 }
 
 // 1 conn is needed to make a session
-func MakeSession(id uint32, valve *Valve, obfs func(*Frame) []byte, deobfs func([]byte) *Frame, obfsedRead func(net.Conn, []byte) (int, error)) *Session {
+func MakeSession(id uint32, valve *Valve, obfs Obfser, deobfs Deobfser, obfsedRead func(net.Conn, []byte) (int, error)) *Session {
 	sesh := &Session{
 		id:           id,
 		obfs:         obfs,

internal/multiplex/stream.go

@@ -98,7 +98,10 @@ func (stream *Stream) Write(in []byte) (n int, err error) {
 		Payload:  in,
 	}
 
-	tlsRecord := stream.session.obfs(f)
+	tlsRecord, err := stream.session.obfs(f)
+	if err != nil {
+		return 0, err
+	}
 	n, err = stream.session.sb.send(tlsRecord)
 	stream.writingM.RUnlock()
 
@@ -132,7 +135,7 @@ func (stream *Stream) Close() error {
 		Closing:  1,
 		Payload:  pad,
 	}
-	tlsRecord := stream.session.obfs(f)
+	tlsRecord, _ := stream.session.obfs(f)
 	stream.session.sb.send(tlsRecord)
 
 	stream.session.delStream(stream.id)

internal/multiplex/switchboard.go

@@ -152,7 +152,11 @@ func (sb *switchboard) deplex(ce *connEnclave) {
 			sb.session.Close()
 			return
 		}
-		frame := sb.session.deobfs(buf[:n])
+		frame, err := sb.session.deobfs(buf[:n])
+		if err != nil {
+			log.Println(err)
+			continue
+		}
 
 		// FIXME: there has been a bug in which a packet has
 		// a seemingly corrupted StreamID (e.g. when the largest streamID is something like 3000

internal/server/usermanager/user.go

@@ -89,7 +89,7 @@ func (u *User) DelSession(sessionID uint32) {
 	u.sessionsM.Unlock()
 }
 
-func (u *User) GetOrCreateSession(sessionID uint32, obfs func(*mux.Frame) []byte, deobfs func([]byte) *mux.Frame, obfsedRead func(net.Conn, []byte) (int, error)) (sesh *mux.Session, existing bool) {
+func (u *User) GetOrCreateSession(sessionID uint32, obfs mux.Obfser, deobfs mux.Deobfser, obfsedRead func(net.Conn, []byte) (int, error)) (sesh *mux.Session, existing bool) {
 	// TODO: session cap
 	u.sessionsM.Lock()
 	if sesh = u.sessions[sessionID]; sesh != nil {