From 50a9633db0a27fa469c37b88d34001c50b9ea7ec Mon Sep 17 00:00:00 2001
From: David Stainton <dstainton415@gmail.com>
Date: Mon, 30 May 2016 20:11:11 +0000
Subject: [PATCH] work-in-progress unit tests for socks proxy chain

---
 socks_server_chain_test.go | 125 ++++++++++++++++++++++++++++++-------
 1 file changed, 104 insertions(+), 21 deletions(-)

diff --git a/socks_server_chain_test.go b/socks_server_chain_test.go
index bc782f3..c506179 100644
--- a/socks_server_chain_test.go
+++ b/socks_server_chain_test.go
@@ -4,11 +4,14 @@ import (
 	"bufio"
 	"bytes"
 	"fmt"
-	"golang.org/x/net/proxy"
+	"io"
 	"net"
 	"strings"
 	"sync"
 	"testing"
+
+	"github.com/subgraph/fw-daemon/socks5"
+	"golang.org/x/net/proxy"
 )
 
 type AccumulatingService struct {
@@ -17,6 +20,7 @@ type AccumulatingService struct {
 	mortalService   *MortalService
 	hasProtocolInfo bool
 	hasAuthenticate bool
+	receivedChan    chan bool
 }
 
 func NewAccumulatingService(net, address string) *AccumulatingService {
@@ -39,53 +43,132 @@ func (a *AccumulatingService) Stop() {
 	a.mortalService.Stop()
 }
 
+func (a *AccumulatingService) WaitUntilReceived() {
+	<-a.receivedChan
+}
+
 func (a *AccumulatingService) SessionWorker(conn net.Conn) error {
 	connReader := bufio.NewReader(conn)
 	for {
 
 		line, err := connReader.ReadBytes('\n')
 		if err != nil {
-			fmt.Println("AccumulatingService read error:", err)
+			panic(fmt.Sprintf("AccumulatingService read error: %s", err))
 		}
 		lineStr := strings.TrimSpace(string(line))
 		a.buffer.WriteString(lineStr + "\n")
+		a.receivedChan <- true
+	}
+	return nil
+}
+
+func fakeSocksSessionWorker(clientConn net.Conn, targetNet, targetAddr string) error {
+	defer clientConn.Close()
+
+	clientAddr := clientConn.RemoteAddr()
+	fmt.Printf("INFO/socks: New connection from: %v\n", clientAddr)
+
+	// Do the SOCKS handshake with the client, and read the command.
+	fmt.Println("meow1")
+	req, err := socks5.Handshake(clientConn)
+	if err != nil {
+		panic(fmt.Sprintf("ERR/socks: Failed SOCKS5 handshake: %v", err))
+	}
+
+	var upstreamConn net.Conn
+	upstreamConn, err = net.Dial(targetNet, targetAddr)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Println("meow2")
+	defer upstreamConn.Close()
+	req.Reply(socks5.ReplySucceeded)
+	fmt.Println("meow3")
+
+	// A upstream connection has been established, push data back and forth
+	// till the session is done.
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	//upstreamConn.Write([]byte("meow 123\r\n"))
+	fmt.Println("meow4")
+	copyLoop := func(dst, src net.Conn) {
+		defer wg.Done()
+		defer dst.Close()
+
+		io.Copy(dst, src)
 	}
+	go copyLoop(upstreamConn, clientConn)
+	go copyLoop(clientConn, upstreamConn)
+
+	fmt.Println("meow5")
+
+	wg.Wait()
+	fmt.Printf("INFO/socks: Closed SOCKS connection from: %v\n", clientAddr)
 	return nil
 }
 
 func TestSocksServerProxyChain(t *testing.T) {
+	// socks client ---> socks chain ---> socks server ---> service
+	socksChainNet := "tcp"
+	socksChainAddr := "127.0.0.1:7750"
+	socksServerNet := "tcp"
+	socksServerAddr := "127.0.0.1:8850"
+	serviceNet := "tcp"
+	serviceAddr := "127.0.0.1:9950"
+
+	fmt.Println("foo1")
+	// setup the service listener
+	service := NewAccumulatingService(serviceNet, serviceAddr)
+	service.Start()
+	defer service.Stop()
+
+	fmt.Println("foo2")
+	// setup the "socks server"
+	session := func(clientConn net.Conn) error {
+		return fakeSocksSessionWorker(clientConn, serviceNet, serviceAddr)
+	}
+	socksService := NewMortalService(socksServerNet, socksServerAddr, session)
+	socksService.Start()
+	defer socksService.Stop()
+
+	fmt.Println("foo3")
+	// setup the SOCKS proxy chain
 	socksConfig := SocksChainConfig{
-		TargetSocksNet:  "tcp",
-		TargetSocksAddr: "127.0.0.1:9050",
-		ListenSocksNet:  "tcp",
-		ListenSocksAddr: "127.0.0.1:8850",
+		TargetSocksNet:  socksServerNet,
+		TargetSocksAddr: socksServerAddr,
+		ListenSocksNet:  socksChainNet,
+		ListenSocksAddr: socksChainAddr,
 	}
 	wg := sync.WaitGroup{}
 	InitSocksListener(&socksConfig, &wg)
 
+	fmt.Println("foo4")
+	// setup the SOCKS client
 	auth := proxy.Auth{
 		User:     "",
 		Password: "",
 	}
 	forward := proxy.NewPerHost(proxy.Direct, proxy.Direct)
-
-	terminatingService := NewAccumulatingService("tcp", "127.0.0.1:1234")
-	terminatingService.Start()
-
-	socksClient, err := proxy.SOCKS5("tcp", "127.0.0.1:8850", &auth, forward)
-	conn, err := socksClient.Dial("tcp", "127.0.0.1:1234")
-
+	socksClient, err := proxy.SOCKS5(socksChainNet, socksChainAddr, &auth, forward)
+	conn, err := socksClient.Dial(serviceNet, serviceAddr)
 	if err != nil {
 		panic(err)
 	}
 
-	rd := bufio.NewReader(conn)
-	line := []byte{}
-	line, err = rd.ReadBytes('\n')
-	if err != nil {
-		panic(err)
-	}
-	fmt.Println("socks client received", string(line))
+	fmt.Println("foo5")
+	conn.Write([]byte("meow 123\r\n"))
+	service.WaitUntilReceived()
+	fmt.Println("DATA RECEIVED", service.buffer.String())
+	fmt.Println("foo6")
+	// read a banner from the service
+	//rd := bufio.NewReader(conn)
+	//line := []byte{}
+	//line, err = rd.ReadBytes('\n')
+	//if err != nil {
+	//	panic(err)
+	//}
+	//fmt.Println("socks client received", string(line))
 
-	wg.Wait()
+	//wg.Wait()
 }