From f76a66fc555150cf752040193fa2d154cbc61b61 Mon Sep 17 00:00:00 2001
From: Jakob Borg <jakob@nym.se>
Date: Thu, 20 Aug 2015 12:59:44 +0200
Subject: [PATCH] Very basic status service

---
 main.go    |  7 +++++++
 session.go | 12 ++++++++++--
 status.go  | 39 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 2 deletions(-)
 create mode 100644 status.go

diff --git a/main.go b/main.go
index b429d94a3..2defed871 100644
--- a/main.go
+++ b/main.go
@@ -32,6 +32,8 @@ var (
 	globalLimitBps  int
 	sessionLimiter  *ratelimit.Bucket
 	globalLimiter   *ratelimit.Bucket
+
+	statusAddr string
 )
 
 func main() {
@@ -47,6 +49,7 @@ func main() {
 	flag.IntVar(&sessionLimitBps, "per-session-rate", sessionLimitBps, "Per session rate limit, in bytes/s")
 	flag.IntVar(&globalLimitBps, "global-rate", globalLimitBps, "Global rate limit, in bytes/s")
 	flag.BoolVar(&debug, "debug", false, "Enable debug output")
+	flag.StringVar(&statusAddr, "status-srv", ":22070", "Listen address for status service (blank to disable)")
 
 	flag.Parse()
 
@@ -99,5 +102,9 @@ func main() {
 
 	go sessionListener(listenSession)
 
+	if statusAddr != "" {
+		go statusService(statusAddr)
+	}
+
 	protocolListener(listenProtocol, tlsCfg)
 }
diff --git a/session.go b/session.go
index c526ed5d3..a189ab403 100644
--- a/session.go
+++ b/session.go
@@ -9,6 +9,7 @@ import (
 	"log"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/juju/ratelimit"
@@ -18,8 +19,10 @@ import (
 )
 
 var (
-	sessionMut = sync.Mutex{}
-	sessions   = make(map[string]*session, 0)
+	sessionMut   = sync.Mutex{}
+	sessions     = make(map[string]*session, 0)
+	numProxies   int64
+	bytesProxied int64
 )
 
 type session struct {
@@ -178,6 +181,9 @@ func (s *session) proxy(c1, c2 net.Conn) error {
 		log.Println("Proxy", c1.RemoteAddr(), "->", c2.RemoteAddr())
 	}
 
+	atomic.AddInt64(&numProxies, 1)
+	defer atomic.AddInt64(&numProxies, -1)
+
 	buf := make([]byte, 65536)
 	for {
 		c1.SetReadDeadline(time.Now().Add(networkTimeout))
@@ -186,6 +192,8 @@ func (s *session) proxy(c1, c2 net.Conn) error {
 			return err
 		}
 
+		atomic.AddInt64(&bytesProxied, int64(n))
+
 		if debug {
 			log.Printf("%d bytes from %s to %s", n, c1.RemoteAddr(), c2.RemoteAddr())
 		}
diff --git a/status.go b/status.go
new file mode 100644
index 000000000..d0391e17d
--- /dev/null
+++ b/status.go
@@ -0,0 +1,39 @@
+package main
+
+import (
+	"encoding/json"
+	"log"
+	"net/http"
+	"runtime"
+	"sync/atomic"
+)
+
+func statusService(addr string) {
+	http.HandleFunc("/status", getStatus)
+	if err := http.ListenAndServe(addr, nil); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func getStatus(w http.ResponseWriter, r *http.Request) {
+	status := make(map[string]interface{})
+
+	sessionMut.Lock()
+	status["numSessions"] = len(sessions)
+	sessionMut.Unlock()
+	status["numProxies"] = atomic.LoadInt64(&numProxies)
+	status["bytesProxied"] = atomic.LoadInt64(&bytesProxied)
+	status["goVersion"] = runtime.Version()
+	status["goOS"] = runtime.GOOS
+	status["goAarch"] = runtime.GOARCH
+	status["goMaxProcs"] = runtime.GOMAXPROCS(-1)
+
+	bs, err := json.MarshalIndent(status, "", "    ")
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.Write(bs)
+}