From 693544659f87f17241112422bacadaa6aee8c8c7 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Wed, 8 Apr 2020 21:37:21 +0100
Subject: [PATCH] Refactor server transport

---
 cmd/ck-server/ck-server.go   |  4 +---
 internal/server/TLS.go       | 10 ++++------
 internal/server/transport.go |  2 --
 internal/server/websocket.go |  4 +---
 internal/util/websocket.go   |  6 ------
 5 files changed, 6 insertions(+), 20 deletions(-)

diff --git a/cmd/ck-server/ck-server.go b/cmd/ck-server/ck-server.go
index 52134cc..484813e 100644
--- a/cmd/ck-server/ck-server.go
+++ b/cmd/ck-server/ck-server.go
@@ -73,7 +73,7 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 
 	var sessionKey [32]byte
 	util.CryptoRandRead(sessionKey[:])
-	obfuscator, err := mux.MakeObfuscator(ci.EncryptionMethod, sessionKey, ci.Transport.HasRecordLayer())
+	obfuscator, err := mux.MakeObfuscator(ci.EncryptionMethod, sessionKey)
 	if err != nil {
 		log.Error(err)
 		goWeb()
@@ -93,7 +93,6 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 		seshConfig := mux.SessionConfig{
 			Obfuscator: obfuscator,
 			Valve:      nil,
-			UnitRead:   ci.Transport.UnitReadFunc(),
 		}
 		sesh := mux.MakeSession(0, seshConfig)
 		sesh.AddConnection(preparedConn)
@@ -125,7 +124,6 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 	seshConfig := mux.SessionConfig{
 		Obfuscator: obfuscator,
 		Valve:      nil,
-		UnitRead:   ci.Transport.UnitReadFunc(),
 		Unordered:  ci.Unordered,
 	}
 	sesh, existing, err := user.GetSession(ci.SessionId, seshConfig)
diff --git a/internal/server/TLS.go b/internal/server/TLS.go
index 7cdcf47..d40ce15 100644
--- a/internal/server/TLS.go
+++ b/internal/server/TLS.go
@@ -15,9 +15,7 @@ type TLS struct{}
 
 var ErrBadClientHello = errors.New("non (or malformed) ClientHello")
 
-func (TLS) String() string                                    { return "TLS" }
-func (TLS) HasRecordLayer() bool                              { return true }
-func (TLS) UnitReadFunc() func(net.Conn, []byte) (int, error) { return util.ReadTLS }
+func (TLS) String() string { return "TLS" }
 
 func (TLS) processFirstPacket(clientHello []byte, privateKey crypto.PrivateKey) (fragments authFragments, respond Responder, err error) {
 	ch, err := parseClientHello(clientHello)
@@ -40,18 +38,18 @@ func (TLS) processFirstPacket(clientHello []byte, privateKey crypto.PrivateKey)
 
 func (TLS) makeResponder(clientHelloSessionId []byte, sharedSecret [32]byte) Responder {
 	respond := func(originalConn net.Conn, sessionKey [32]byte) (preparedConn net.Conn, err error) {
-		preparedConn = originalConn
 		reply, err := composeReply(clientHelloSessionId, sharedSecret, sessionKey)
 		if err != nil {
 			err = fmt.Errorf("failed to compose TLS reply: %v", err)
 			return
 		}
-		_, err = preparedConn.Write(reply)
+		_, err = originalConn.Write(reply)
 		if err != nil {
 			err = fmt.Errorf("failed to write TLS reply: %v", err)
-			go preparedConn.Close()
+			go originalConn.Close()
 			return
 		}
+		preparedConn = &util.TLSConn{Conn: originalConn}
 		return
 	}
 	return respond
diff --git a/internal/server/transport.go b/internal/server/transport.go
index 375a016..0b2a369 100644
--- a/internal/server/transport.go
+++ b/internal/server/transport.go
@@ -8,8 +8,6 @@ import (
 
 type Responder = func(originalConn net.Conn, sessionKey [32]byte) (preparedConn net.Conn, err error)
 type Transport interface {
-	HasRecordLayer() bool
-	UnitReadFunc() func(net.Conn, []byte) (int, error)
 	processFirstPacket(reqPacket []byte, privateKey crypto.PrivateKey) (authFragments, Responder, error)
 }
 
diff --git a/internal/server/websocket.go b/internal/server/websocket.go
index a05bcb2..0a910c5 100644
--- a/internal/server/websocket.go
+++ b/internal/server/websocket.go
@@ -15,9 +15,7 @@ import (
 
 type WebSocket struct{}
 
-func (WebSocket) String() string                                    { return "WebSocket" }
-func (WebSocket) HasRecordLayer() bool                              { return false }
-func (WebSocket) UnitReadFunc() func(net.Conn, []byte) (int, error) { return util.ReadWebSocket }
+func (WebSocket) String() string { return "WebSocket" }
 
 func (WebSocket) processFirstPacket(reqPacket []byte, privateKey crypto.PrivateKey) (fragments authFragments, respond Responder, err error) {
 	var req *http.Request
diff --git a/internal/util/websocket.go b/internal/util/websocket.go
index 3ea49db..574f51c 100644
--- a/internal/util/websocket.go
+++ b/internal/util/websocket.go
@@ -4,7 +4,6 @@ import (
 	"errors"
 	"github.com/gorilla/websocket"
 	"io"
-	"net"
 	"sync"
 	"time"
 )
@@ -75,8 +74,3 @@ func (ws *WebSocketConn) SetDeadline(t time.Time) error {
 	}
 	return nil
 }
-
-// ws unit reader
-func ReadWebSocket(conn net.Conn, buffer []byte) (n int, err error) {
-	return conn.Read(buffer)
-}