Network proxy in 20 lines

Authentication, load balancing… You need a proxy ? Let’s do it in an idiomatic way.
It’s surprisingly easy !

Thumbnail

Needs

We need to handle 2 connections : the User and the Server.

Story :
1- The User connects to the Proxy
2- The Proxy connects to the Server
3- The Proxy sent messages from the Client to the Server, and from the Server to the Client (transparently)

We don’t speak about the specific treatment of the proxy (authentication, load balancing …) you wish to do.

Problematic

Once you get your 2 connections, you naturally want to do something like that :

 
func startProxy(client net.Conn, server net.Conn) {
  go relay(client, server)
  go relay(server, client)
}
 
func relay(src net.Conn, dst net.Conn) {
  for {
    var buff []byte
    _, err := src.Read(buff)
    if err != nil {
      return
    }
    dst.Write(buff)
  }
}

What is untoward here :
When the client disconnects or even if the server disconnects (crash), it’s not possible to close both connections.
Example :
1. The client disconnects. The goroutine reading the client exits as expected.
2. The goroutines reading the server is blocked in Read.

The Gopher first reaction here is to add a “Cancellation” channel and to add a Deadline at each Read.

func relay(src net.Conn, dst net.Conn, cancellation chan bool) {
  for {
    src.SetReadDeadline(time.Now().Add(time.Second*10))
 
    var buff []byte
    _, err := src.Read(buff)
 
    if err == err.Timeout() { // True when read exit because of timeout
      return
    } else if err != nil { // Something wrong
      cancellation <- true
    }
    select {
    case <- cancellation:
        return
    default:
        // Do nothing. Just for non-blocking channel read
    }
    dst.Write(buff)
  }
}
Blocking read/write and deadlines
By default, the Read and Write on net.Conn are blocking.

But you can change this with method “SetReadDeadline” and “SetWriteDeadline”. (Or just “SetDeadline” for both in once)

You specify a moment (with time.Time object) like “time.Now().Add(time.Second*10)” and you got 10 seconds of wait before the Read/Write returns if nothing is Read/Written.

This said, there is a way to do it a lot more simple.

Solution

2 tips are enough to do it better :
– Do you remember “io.Copy()” function ? Give a look.
– Do you know that a connection can be manually closed by calling the method “Closed” of the net.Conn interface ?

Here we are :

func startProxy(client net.Conn, server net.Conn) {
  stop := make(chan bool)
 
  go relay(client, server, stop)
  go relay(server, client, stop)
 
  select {
    case <-stop:
      return
  }
}
 
func relay(src net.Conn, dst net.Conn, stop chan bool) {
	io.Copy(dst, src)
	dst.Close()
	src.Close()
	stop <- true
	return
}

The important things here :
1. “io.Copy” returns only if something wrong happend (including disconnection)
2. About “Close()”, the doc says : “Any blocked Read or Write operations will be unblocked and return errors.”

Did you know ?
The net.Conn interface allows you to handle Websocket connections too !
This solution works for both.

It can even be useful to convert websocket to socket or the opposite.

A proxy in Golang can be very performing.
It was also an opportunity to talk about the possibilities of “conn” interface in the net package.

Entrepreneur – Cofounder at Golem.ai (Paris, France)

I enjoy sharing Golang interesting patterns, experiments and tips.

1 thought on “Network proxy in 20 lines

Leave a Reply

Your email address will not be published. Required fields are marked *