c93614 · September 4, 2019 03:19 · Jun 16, 2013
diff --git a/balance.go b/balance.go
@@ -0,0 +1,80 @@
+package main
+
+import (
+    "flag"
+    "io"
+    "log"
+    "net"
+    "strings"
+)
+
+type Backends struct {
+    servers []string
+    n       int
+}
+
+func (b *Backends) Choose() string {
+    idx := b.n % len(b.servers)
+    b.n++
+    return b.servers[idx]
+}
+
+func (b *Backends) String() string {
+    return strings.Join(b.servers, ", ")
+}
+
+var (
+    bind     = flag.String("bind", "", "The address to bind on")
+    balance  = flag.String("balance", "", "The backend servers to balance connections across, separated by commas")
+    backends *Backends
+)
+
+func init() {
+    flag.Parse()
+
+    if *bind == "" {
+        log.Fatalln("specify the address to listen on with -bind")
+    }
+
+    servers := strings.Split(*balance, ",")
+    if len(servers) == 1 && servers[0] == "" {
+        log.Fatalln("please specify backend servers with -backends")
+    }
+
+    backends = &Backends{servers: servers}
+}
+
+func copy(wc io.WriteCloser, r io.Reader) {
+    defer wc.Close()
+    io.Copy(wc, r)
+}
+
+func handleConnection(us net.Conn, server string) {
+    ds, err := net.Dial("tcp", server)
+    if err != nil {
+        us.Close()
+        log.Printf("failed to dial %s: %s", server, err)
+        return
+    }
+
+    go copy(ds, us)
+    go copy(us, ds)
+}
+
+func main() {
+    ln, err := net.Listen("tcp", *bind)
+    if err != nil {
+        log.Fatalf("failed to bind: %s", err)
+    }
+
+    log.Printf("listening on %s, balancing %s", *bind, backends)
+
+    for {
+        conn, err := ln.Accept()
+        if err != nil {
+            log.Printf("failed to accept: %s", err)
+            continue
+        }
+        go handleConnection(conn, backends.Choose())
+    }
+}