Skip to content

Instantly share code, notes, and snippets.

@olekukonko
Last active December 25, 2015 18:19
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save olekukonko/7019992 to your computer and use it in GitHub Desktop.
Save olekukonko/7019992 to your computer and use it in GitHub Desktop.
package main
import (
"flag"
"fmt"
"log"
"net"
)
var fromHost = flag.String("from", "localhost:80", "The proxy server's host.")
var toHost = flag.String("to", "localhost:8000", "The host that the proxy " +
" server should forward requests to.")
var maxConnections = flag.Int("c", 25, "The maximum number of active " +
"connection at any given time.")
var maxWaitingConnections = flag.Int("cw", 10000, "The maximum number of " +
"connections that can be waiting to be served.")
func main() {
// Parse the command-line arguments.
flag.Parse()
fmt.Printf("Proxying %s->%s.\r\n", *fromHost, *toHost)
// Set up our listening server
server, err := net.Listen("tcp", *fromHost)
// If any error occurs while setting up our listening server, error out.
if err != nil {
log.Fatal(err)
}
// The channel of connections which are waiting to be processed.
waiting := make(chan net.Conn, *maxWaitingConnections)
// The booleans representing the free active connection spaces.
spaces := make(chan bool, *maxConnections)
// Initialize the spaces
for i := 0; i < *maxConnections; i++ {
spaces <- true
}
// Start the connection matcher.
go matchConnections(waiting, spaces)
// Loop indefinitely, accepting connections and handling them.
for {
connection, err := server.Accept()
if err != nil {
// Log the error.
log.Print(err)
} else {
// Create a goroutine to handle the conn
log.Printf("Received connection from %s.\r\n",
connection.RemoteAddr())
waiting <- connection
}
}
}
func matchConnections(waiting chan net.Conn, spaces chan bool) {
// Iterate over each connection in the waiting channel
for connection := range waiting {
// Block until we have a space.
<-spaces
// Create a new goroutine which will call the connection handler and
// then free up the space.
go func(connection net.Conn) {
handleConnection(connection)
spaces <- true
log.Printf("Closed connection from %s.\r\n", connection.RemoteAddr())
}(connection)
}
}
func handleConnection(connection net.Conn) {
// Always close our connection.
defer connection.Close()
// Try to connect to remote server.
remote, err := net.Dial("tcp", *toHost)
if err != nil {
// Exit out when an error occurs
log.Print(err)
return
}
defer remote.Close()
// Create our channel which waits for completion, and our two channels to
// signal that a goroutine is done.
complete := make(chan bool, 2)
ch1 := make(chan bool, 1)
ch2 := make(chan bool, 1)
go copyContent(connection, remote, complete, ch1, ch2)
go copyContent(remote, connection, complete, ch2, ch1)
// Block until we've completed both goroutines!
<- complete
<- complete
}
func copyContent(from net.Conn, to net.Conn, complete chan bool, done chan bool, otherDone chan bool) {
var err error = nil
var bytes []byte = make([]byte, 256)
var read int = 0
for {
select {
// If we received a done message from the other goroutine, we exit.
case <- otherDone:
complete <- true
return
default:
// Read data from the source connection.
read, err = from.Read(bytes)
// If any errors occured, write to complete as we are done (one of the
// connections closed.)
if err != nil {
complete <- true
done <- true
return
}
// Write data to the destination.
_, err = to.Write(bytes[:read])
// Same error checking.
if err != nil {
complete <- true
done <- true
return
}
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment