proxy.go

package main

import (
	"fmt"
	"net"
	"bufio"
	"net/http"
	"net/http/httputil"
	"io"
	"time"
	"strings"
	"runtime"
	"strconv"
)

func respondConnectionEstablished(socket net.Conn) (n int, err error) {
	return socket.Write([]byte("HTTP/1.0 200 Connection established\r\n\r\n"))
}

func respondOk(socket net.Conn) (n int, err error) {
	return socket.Write([]byte("HTTP/1.1 200 OK\r\n\r\n"))
}

func acceptSocketChannel(listener net.Listener) chan net.Conn {
	channel := make(chan net.Conn)
	go func() {
		for {
			socket, err := listener.Accept()
			if err != nil {
				fmt.Printf("Could not accept socket: " + err.Error())
				continue
			}

			channel <- socket
		}
	}()
	return channel
}

func proxyHttpReuqest(socket net.Conn, request * http.Request) {
	// For some reason we can't have the RequestURI when using http.Client{}
	request.RequestURI = ""

	client := http.Client{}
	response, err := client.Do(request)
	if err != nil {
		fmt.Println("Error sending request upstream: " + err.Error())
		return
	}
	defer response.Body.Close()

	dump, err := httputil.DumpResponse(response, true)
	if err != nil {
		fmt.Println("Error dumping upstream response to bytes: " + err.Error())
		return
	}

	socket.Write(dump)
}

func handleSocket(socket net.Conn) {
	defer socket.Close()
	reader := bufio.NewReader(socket)
	// TODO: Do we need to close this request body?
	request, err := http.ReadRequest(reader)
	if err != nil {
		fmt.Printf("Could not read request from socket: " + err.Error())
		return;
	}

	fmt.Println("Upstream host: " + request.URL.Host)
	fmt.Println("Upstream uri: " + request.URL.RequestURI())

	if request.Method == "CONNECT" {
		fmt.Println("SSL socket transport")
		fmt.Println("Found SHTTP proxy URL: " + request.URL.String())
		proxyRawBytes(socket, request)
	} else {
		fmt.Println("Found proxy URL: " + request.Method)
		fmt.Println("Found HTTP proxy URL: " + request.URL.String())
		proxyHttpReuqest(socket, request)
	}
}

func proxyRawBytes(socket net.Conn, request * http.Request) {
	host := request.URL.Host

	fmt.Println("Opening TCP connection to host: " + host)

	upstreamSocket, err := net.Dial("tcp", host)
	if err != nil {
		fmt.Println("Error establishing socket connection upstream: " + err.Error())
	}
	defer upstreamSocket.Close()

	_, err = respondConnectionEstablished(socket)
	if err != nil {
		fmt.Println("Could not send Continue response to client: " + err.Error())
	}

	// At this point, the sockets should be healthy. Let's relay.
	signal := make(chan int)

	go spawnRelayWithSignal(socket, upstreamSocket, "client->upstream", signal)
	go spawnRelayWithSignal(upstreamSocket, socket, "upstream->client", signal)

	// Wait for reader/ writer chans to finish
	// TODO: I think this will leave a dangling goroutine?
	<-signal
	fmt.Println("Closing connection")
}

func spawnRelayWithSignal(readFrom net.Conn, writeTo net.Conn, tag string, signal chan int) {
	readFrom.SetReadDeadline(time.Now().Add(60 * time.Second))
	_, err := io.Copy(writeTo, readFrom)

	if err != nil {
		errString := err.Error()

		if !strings.Contains(errString, "use of closed network connection") {
			fmt.Println("Error copying - " + tag + ": " + err.Error())
		}
	}

	signal <- 0
}

func periodicStatPrinter() {
	for {
		fmt.Println("GoRoutines: " + strconv.Itoa(runtime.NumGoroutine()))
		time.Sleep(10 * time.Second)
	}
}

func main() {
	server, err := net.Listen("tcp", ":25000")
	if err != nil {
		panic("Could not start server: " + err.Error())
	}

	fmt.Println("Started server on port 25000")

	go periodicStatPrinter()

	acceptSocketChannel := acceptSocketChannel(server)
	for {
		go handleSocket(<-acceptSocketChannel)
	}
}