From 94258abb365aefd04ea652a3e024a807d7af14e1 Mon Sep 17 00:00:00 2001
From: David Stainton <dstainton415@gmail.com>
Date: Mon, 30 May 2016 20:48:24 +0000
Subject: [PATCH] Make the test receive a banner through the proxy chain from
 the service

---
 socks_server_chain_test.go | 44 ++++++++++++++------------------------
 1 file changed, 16 insertions(+), 28 deletions(-)

diff --git a/socks_server_chain_test.go b/socks_server_chain_test.go
index c506179..9656607 100644
--- a/socks_server_chain_test.go
+++ b/socks_server_chain_test.go
@@ -16,6 +16,7 @@ import (
 
 type AccumulatingService struct {
 	net, address    string
+	banner          string
 	buffer          bytes.Buffer
 	mortalService   *MortalService
 	hasProtocolInfo bool
@@ -23,10 +24,11 @@ type AccumulatingService struct {
 	receivedChan    chan bool
 }
 
-func NewAccumulatingService(net, address string) *AccumulatingService {
+func NewAccumulatingService(net, address, banner string) *AccumulatingService {
 	l := AccumulatingService{
 		net:             net,
 		address:         address,
+		banner:          banner,
 		hasProtocolInfo: true,
 		hasAuthenticate: true,
 	}
@@ -49,8 +51,8 @@ func (a *AccumulatingService) WaitUntilReceived() {
 
 func (a *AccumulatingService) SessionWorker(conn net.Conn) error {
 	connReader := bufio.NewReader(conn)
+	conn.Write([]byte(a.banner))
 	for {
-
 		line, err := connReader.ReadBytes('\n')
 		if err != nil {
 			panic(fmt.Sprintf("AccumulatingService read error: %s", err))
@@ -69,7 +71,6 @@ func fakeSocksSessionWorker(clientConn net.Conn, targetNet, targetAddr string) e
 	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))
@@ -80,18 +81,15 @@ func fakeSocksSessionWorker(clientConn net.Conn, targetNet, targetAddr string) e
 	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")
+	//clientConn.Write([]byte("meow 123\r\n"))
 	copyLoop := func(dst, src net.Conn) {
 		defer wg.Done()
 		defer dst.Close()
@@ -101,8 +99,6 @@ func fakeSocksSessionWorker(clientConn net.Conn, targetNet, targetAddr string) e
 	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
@@ -117,13 +113,12 @@ func TestSocksServerProxyChain(t *testing.T) {
 	serviceNet := "tcp"
 	serviceAddr := "127.0.0.1:9950"
 
-	fmt.Println("foo1")
+	banner := "meow 123\r\n"
 	// setup the service listener
-	service := NewAccumulatingService(serviceNet, serviceAddr)
+	service := NewAccumulatingService(serviceNet, serviceAddr, banner)
 	service.Start()
 	defer service.Stop()
 
-	fmt.Println("foo2")
 	// setup the "socks server"
 	session := func(clientConn net.Conn) error {
 		return fakeSocksSessionWorker(clientConn, serviceNet, serviceAddr)
@@ -132,7 +127,6 @@ func TestSocksServerProxyChain(t *testing.T) {
 	socksService.Start()
 	defer socksService.Stop()
 
-	fmt.Println("foo3")
 	// setup the SOCKS proxy chain
 	socksConfig := SocksChainConfig{
 		TargetSocksNet:  socksServerNet,
@@ -143,7 +137,6 @@ func TestSocksServerProxyChain(t *testing.T) {
 	wg := sync.WaitGroup{}
 	InitSocksListener(&socksConfig, &wg)
 
-	fmt.Println("foo4")
 	// setup the SOCKS client
 	auth := proxy.Auth{
 		User:     "",
@@ -156,19 +149,14 @@ func TestSocksServerProxyChain(t *testing.T) {
 		panic(err)
 	}
 
-	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()
+	rd := bufio.NewReader(conn)
+	line := []byte{}
+	line, err = rd.ReadBytes('\n')
+	if err != nil {
+		panic(err)
+	}
+	if string(line) != banner {
+		t.Errorf("Did not receive expected banner. Got %s, wanted %s\n", string(line), banner)
+	}
 }