2 files changed, 17 insertions, 0 deletions
diff --git a/log.go b/log.go
index f2305ce..54489f3 100644
--- a/log.go
+++ b/log.go
@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"os"
 	"path"
+	"strings"
 	"sync"
 	"time"
 )
@@ -22,6 +23,7 @@ type Logger struct {
 	HashUserAgent bool
 	HashReferer   bool
 	HashSalt      string
+	ProxyCount    int
 	logFile       *os.File
 	encoder       *json.Encoder
 	lastDate      string
@@ -52,6 +54,19 @@ func (l *Logger) Log(entry LogEntry) {
 
 func (l *Logger) LogUpload(req *http.Request, res result) {
 	host, _, _ := net.SplitHostPort(req.RemoteAddr)
+	if l.ProxyCount > 0 {
+		ffs := strings.Split(req.Header.Get("X-Forwarded-For"), ",")
+		if len(ffs) < l.ProxyCount {
+			ri := req.Header.Get("X-Real-IP")
+			if ri != "" {
+				host = ri
+			} else if len(ffs) > 0 {
+				host = ffs[len(ffs)-1]
+			}
+		} else {
+			host = ffs[len(ffs)-l.ProxyCount]
+		}
+	}
 	l.logUpload(
 		host,               // ip
 		req.UserAgent(),    // userAgent
diff --git a/main.go b/main.go
index e8ff33b..7953398 100644
--- a/main.go
+++ b/main.go
@@ -85,6 +85,7 @@ func main() {
 	logReferer := flag.Bool("log-referer", false, "log Referer headers")
 	logRefererHash := flag.Bool("log-referer-hash", false, "log hashed Referer headers")
 	logHashSalt := flag.String("log-hash-salt", "", "salt to use for hashed log entries")
+	proxyCount := flag.Int("proxy-count", 0, "count of trusted reverse proxies")
 
 	flag.Parse()
 
@@ -115,6 +116,7 @@ func main() {
 		DefaultLogger.HashUserAgent = *logUAHash
 		DefaultLogger.HashReferer = *logRefererHash
 		DefaultLogger.HashSalt = *logHashSalt
+		DefaultLogger.ProxyCount = *proxyCount
 	}
 
 	http.HandleFunc("/upload.php", handleUpload)