refactor

1 files changed, 40 insertions, 13 deletions

diff --git a/client/main.go b/client/main.go
index 1c866ae..b536a81 100644
--- a/client/main.go
+++ b/client/main.go
@@ -46,30 +46,57 @@ func main() {
 		debug.SetOutput(os.Stderr)
 	}
 
-	dialer, err := obfsvpn.NewDialerFromCert(obfs4Cert)
-	if err != nil {
-		logger.Fatalf("cannot get dialer: %v", err)
-	}
-
+	kcpTransport := false
 	// TODO make this configurable via a Config struct
 	// TODO make sure we're disabling all the crypto options for KCP
 	if os.Getenv("KCP") == "1" {
-		dialer.DialFunc = func(network, address string) (net.Conn, error) {
-			log.Printf("Dialing kcp://%s\n", address)
-			return kcp.Dial(address)
-		}
+		kcpTransport = true
 	}
 
 	socksAddr := net.JoinHostPort(socksHost, socksPort)
+	client := NewClient(kcpTransport, socksAddr, obfs4Cert)
+	client.Start()
+}
+
+type Client struct {
+	kcp       bool
+	socksAddr string
+	obfs4Cert string
+}
 
+func NewClient(kcp bool, socksAddr, obfs4Cert string) *Client {
+	return &Client{
+		kcp:       kcp,
+		socksAddr: socksAddr,
+		obfs4Cert: obfs4Cert,
+	}
+}
+
+func (c *Client) Start() bool {
 	server := &socks5.Server{
-		Addr:   socksAddr,
+		Addr:   c.socksAddr,
 		BindIP: "127.0.0.1",
-		Dial:   dialer.Dial,
 	}
 
-	fmt.Printf("[+] Starting socks5 proxy at %s\n", socksAddr)
+	dialer, err := obfsvpn.NewDialerFromCert(c.obfs4Cert)
+	if err != nil {
+		log.Printf("Error getting dialer: %v\n", err)
+		return false
+	}
+
+	if c.kcp {
+		dialer.DialFunc = func(network, address string) (net.Conn, error) {
+			log.Printf("Dialing kcp://%s\n", address)
+			return kcp.Dial(address)
+		}
+	}
+
+	server.Dial = dialer.Dial
+
+	fmt.Printf("[+] Starting socks5 proxy at %s\n", c.socksAddr)
 	if err := server.ListenAndServe(); err != nil {
-		panic(err)
+		log.Printf("error while listening: %v\n", err)
+		return false
 	}
+	return true
 }