golang websocket to tcp forwarder

I’ve been writing a simple websocket to tcp forwarder in golang for a project I’m working on.

I previously showed a simple echo example using Gorilla websockets. The echoing example blocks on the sockets until data is available, however that isn’t really going to work here. At first I tried to use the SetReadDeadLine, mirroring non-blocking IO in C. While this worked, it wasn’t really an ideal solution. What I ended up with was my first use of go routines! Basically breaking the code into two threads each of which blocks on reads from the 2 sockets. The code is below for reference (it currently forwards traffic to/from localhost, but the first 12bytes sent from the client should be the destination host port, I’ll update this post then that codes in place):

package main
 
import (
    "github.com/gorilla/websocket"
    "net/http"
    "net"
    "fmt"
    "io"
)

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
}


      
func print_binary(s []byte) {
  fmt.Printf("print b:");
  for n := 0;n < len(s);n++ {
    fmt.Printf("%d,",s[n]);
  }
  fmt.Printf("\n");
}

func address_decode(address_bin []byte) (string,string) {
  
  var host string = "127.0.0.1"
  var port string = "22";

  return host,port
}
 

func forwardtcp(wsconn *websocket.Conn,conn net.Conn) {

  for {
    // Receive and forward pending data from tcp socket to web socket
    tcpbuffer := make([]byte, 1024)

    n,err := conn.Read(tcpbuffer)
    if err == io.EOF { fmt.Printf("TCP Read failed"); break; }
    if err == nil {
      fmt.Printf("Forwarding from tcp to ws: %d bytes: %s\n",n,tcpbuffer)
      print_binary(tcpbuffer)
      wsconn.WriteMessage(websocket.BinaryMessage,tcpbuffer[:n])
    }
  }
}

func forwardws (wsconn *websocket.Conn,conn net.Conn) {

 for {
    // Send pending data to tcp socket
    n,buffer,err := wsconn.ReadMessage()
    if err == io.EOF { fmt.Printf("WS Read Failed"); break; }
    if err == nil {
      s := string(buffer[:len(buffer)])
      fmt.Printf("Received (from ws) forwarding to tcp: %d bytes: %s %d\n",len(buffer),s,n)
      print_binary(buffer)
      conn.Write(buffer)
    }
  }
}

func wsProxyHandler(w http.ResponseWriter, r *http.Request) {

  wsconn, err := upgrader.Upgrade(w, r, nil)

  if err != nil {
    //log.Println(err)
    return
  }

  // get connection address and port
  address := make([]byte, 16)

  n,address,err := wsconn.ReadMessage()
  if err != nil {
    fmt.Printf("address read error");
    fmt.Printf("read %d bytes",n);  
  }

  print_binary(address)

  host, port := address_decode(address)

  conn, err := net.Dial("tcp", host + ":" + port)
  if err != nil {
	// handle error
  }

  go forwardtcp(wsconn,conn)
  go forwardws(wsconn,conn)

  fmt.Printf("websocket closed");
}
 
func main() {
  http.HandleFunc("/echo", wsProxyHandler)
  http.Handle("/", http.FileServer(http.Dir(".")))
  err := http.ListenAndServe(":8080", nil)
  if err != nil {
    panic("Error: " + err.Error())
  }
}