From c19c43f6e87b41b40408c0f483def8649a22af77 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Mon, 12 Aug 2019 23:13:13 +0100
Subject: [PATCH] Server side UDP to proxy server

---
 cmd/ck-server/ck-server.go    |  8 +++++--
 internal/multiplex/session.go |  1 +
 internal/server/state.go      | 42 ++++++++++++++++++++++++++++-------
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/cmd/ck-server/ck-server.go b/cmd/ck-server/ck-server.go
index e3f205b..1acfc0b 100644
--- a/cmd/ck-server/ck-server.go
+++ b/cmd/ck-server/ck-server.go
@@ -167,7 +167,8 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 				continue
 			}
 		}
-		localConn, err := net.Dial("tcp", sta.ProxyBook[ci.ProxyMethod])
+		proxyAddr := sta.ProxyBook[ci.ProxyMethod]
+		localConn, err := net.Dial(proxyAddr.Network(), proxyAddr.String())
 		if err != nil {
 			log.Errorf("Failed to connect to %v: %v", ci.ProxyMethod, err)
 			user.DeleteSession(ci.SessionId, "Failed to connect to proxy server")
@@ -254,7 +255,10 @@ func main() {
 		if net.ParseIP(ssLocalHost).To4() == nil {
 			ssLocalHost = "[" + ssLocalHost + "]"
 		}
-		sta.ProxyBook["shadowsocks"] = ssLocalHost + ":" + ssLocalPort
+		sta.ProxyBook["shadowsocks"], err = net.ResolveTCPAddr("tcp", ssLocalHost+":"+ssLocalPort)
+		if err != nil {
+			log.Fatal(err)
+		}
 	}
 
 	listen := func(addr, port string) {
diff --git a/internal/multiplex/session.go b/internal/multiplex/session.go
index cf4b98e..3a5ca47 100644
--- a/internal/multiplex/session.go
+++ b/internal/multiplex/session.go
@@ -80,6 +80,7 @@ func MakeSession(id uint32, config *SessionConfig) *Session {
 		Valve: config.Valve,
 	}
 	if config.Unordered {
+		log.Debug("Connection is unordered")
 		sbConfig.strategy = UNIFORM_SPREAD
 	} else {
 		sbConfig.strategy = FIXED_CONN_MAPPING
diff --git a/internal/server/state.go b/internal/server/state.go
index 0af7d4f..318a02c 100644
--- a/internal/server/state.go
+++ b/internal/server/state.go
@@ -5,8 +5,11 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"github.com/cbeuw/Cloak/internal/server/usermanager"
 	"io/ioutil"
+	"net"
+	"strings"
 	"sync"
 	"time"
 
@@ -14,7 +17,7 @@ import (
 )
 
 type rawConfig struct {
-	ProxyBook    map[string]string
+	ProxyBook    map[string][]string
 	BypassUID    [][]byte
 	RedirAddr    string
 	PrivateKey   string
@@ -25,7 +28,7 @@ type rawConfig struct {
 
 // State type stores the global state of the program
 type State struct {
-	ProxyBook map[string]string
+	ProxyBook map[string]net.Addr
 
 	BindHost string
 	BindPort string
@@ -47,12 +50,13 @@ type State struct {
 
 func InitState(bindHost, bindPort string, nowFunc func() time.Time) (*State, error) {
 	ret := &State{
-		BindHost:  bindHost,
-		BindPort:  bindPort,
-		Now:       nowFunc,
-		BypassUID: make(map[[16]byte]struct{}),
+		BindHost:   bindHost,
+		BindPort:   bindPort,
+		Now:        nowFunc,
+		BypassUID:  make(map[[16]byte]struct{}),
+		ProxyBook:  map[string]net.Addr{},
+		usedRandom: map[[32]byte]int64{},
 	}
-	ret.usedRandom = make(map[[32]byte]int64)
 	go ret.UsedRandomCleaner()
 	return ret, nil
 }
@@ -88,7 +92,29 @@ func (sta *State) ParseConfig(conf string) (err error) {
 	}
 
 	sta.RedirAddr = preParse.RedirAddr
-	sta.ProxyBook = preParse.ProxyBook
+
+	for name, pair := range preParse.ProxyBook {
+		if len(pair) != 2 {
+			return fmt.Errorf("invalid protocol and address pair for %v: %v", name, pair)
+		}
+		network := strings.ToLower(pair[0])
+		switch network {
+		case "tcp":
+			addr, err := net.ResolveTCPAddr("tcp", pair[1])
+			if err != nil {
+				return err
+			}
+			sta.ProxyBook[name] = addr
+			continue
+		case "udp":
+			addr, err := net.ResolveUDPAddr("udp", pair[1])
+			if err != nil {
+				return err
+			}
+			sta.ProxyBook[name] = addr
+			continue
+		}
+	}
 
 	pvBytes, err := base64.StdEncoding.DecodeString(preParse.PrivateKey)
 	if err != nil {