enh: get client real ip

taosdata · Apr 8, 2024 · 283f89e · 283f89e
1 parent 2d055ac
commit 283f89e
Show file tree

Hide file tree

Showing 13 changed files with 139 additions and 84 deletions.
diff --git a/controller/rest/restful.go b/controller/rest/restful.go
@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"net"
 	"net/http"
 	"strconv"
 	"strings"
@@ -29,6 +28,7 @@ import (
 	"github.com/taosdata/taosadapter/v3/tools"
 	"github.com/taosdata/taosadapter/v3/tools/csv"
 	"github.com/taosdata/taosadapter/v3/tools/ctools"
+	"github.com/taosdata/taosadapter/v3/tools/iptool"
 	"github.com/taosdata/taosadapter/v3/tools/jsonbuilder"
 	"github.com/taosdata/taosadapter/v3/tools/layout"
 	"github.com/taosdata/taosadapter/v3/tools/pool"
@@ -167,7 +167,7 @@ func DoQuery(c *gin.Context, db string, timeFunc ctools.FormatTimeFunc, reqID in
 	if isDebug {
 		s = time.Now()
 	}
-	taosConnect, err := commonpool.GetConnection(user, password, net.ParseIP(c.RemoteIP()))
+	taosConnect, err := commonpool.GetConnection(user, password, iptool.GetRealIP(c.Request))
 	if isDebug {
 		logger.Debugln("connect server cost:", time.Since(s))
 	}
@@ -500,13 +500,14 @@ func (ctl *Restful) upload(c *gin.Context) {
 	user := c.MustGet(UserKey).(string)
 	password := c.MustGet(PasswordKey).(string)
 	s := log.GetLogNow(isDebug)
-	taosConnect, err := commonpool.GetConnection(user, password, net.ParseIP(c.RemoteIP()))
+	ip := iptool.GetRealIP(c.Request)
+	taosConnect, err := commonpool.GetConnection(user, password, ip)
 	if isDebug {
 		logger.Debugln("connect cost:", log.GetLogDuration(isDebug, s))
 	}
 
 	if err != nil {
-		logger.WithError(err).Error("connect server error")
+		logger.WithField("ip", ip).WithError(err).Error("connect server error")
 		var tError *tErrors.TaosError
 		if errors.Is(err, commonpool.ErrWhitelistForbidden) {
 			ForbiddenResponse(c, commonpool.ErrWhitelistForbidden.Error())
@@ -674,7 +675,7 @@ func (ctl *Restful) des(c *gin.Context) {
 		BadRequestResponse(c, httperror.HTTP_GEN_TAOSD_TOKEN_ERR)
 		return
 	}
-	conn, err := commonpool.GetConnection(user, password, net.ParseIP(c.RemoteIP()))
+	conn, err := commonpool.GetConnection(user, password, iptool.GetRealIP(c.Request))
 	if err != nil {
 		if errors.Is(err, commonpool.ErrWhitelistForbidden) {
 			ForbiddenResponse(c, commonpool.ErrWhitelistForbidden.Error())

diff --git a/controller/rest/table_vgid.go b/controller/rest/table_vgid.go
@@ -2,7 +2,6 @@ package rest
 
 import (
 	"errors"
-	"net"
 	"net/http"
 
 	"github.com/gin-gonic/gin"
@@ -11,6 +10,7 @@ import (
 	"github.com/taosdata/taosadapter/v3/db/commonpool"
 	"github.com/taosdata/taosadapter/v3/log"
 	"github.com/taosdata/taosadapter/v3/thread"
+	"github.com/taosdata/taosadapter/v3/tools/iptool"
 )
 
 func (ctl *Restful) tableVgID(c *gin.Context) {
@@ -30,7 +30,7 @@ func (ctl *Restful) tableVgID(c *gin.Context) {
 	password := c.MustGet(PasswordKey).(string)
 	isDebug := log.IsDebug()
 	s := log.GetLogNow(isDebug)
-	taosConn, err := commonpool.GetConnection(user, password, net.ParseIP(c.RemoteIP()))
+	taosConn, err := commonpool.GetConnection(user, password, iptool.GetRealIP(c.Request))
 	logger.Debugln("connect cost:", log.GetLogDuration(isDebug, s))
 	if err != nil {
 		logger.WithError(err).Error("connect server error")

diff --git a/controller/ws/query/ws.go b/controller/ws/query/ws.go
@@ -6,7 +6,6 @@ import (
 	"context"
 	"encoding/json"
 	"net"
-	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -27,6 +26,7 @@ import (
 	"github.com/taosdata/taosadapter/v3/monitor"
 	"github.com/taosdata/taosadapter/v3/thread"
 	"github.com/taosdata/taosadapter/v3/tools"
+	"github.com/taosdata/taosadapter/v3/tools/iptool"
 	"github.com/taosdata/taosadapter/v3/tools/jsontype"
 )
 
@@ -220,8 +220,7 @@ type Taos struct {
 }
 
 func NewTaos(session *melody.Session) *Taos {
-	host, _, _ := net.SplitHostPort(strings.TrimSpace(session.Request.RemoteAddr))
-	ipAddr := net.ParseIP(host)
+	ipAddr := iptool.GetRealIP(session.Request)
 	return &Taos{
 		Results:             list.New(),
 		exit:                make(chan struct{}, 1),
@@ -246,7 +245,7 @@ func (t *Taos) waitSignal() {
 				return
 			}
 			logger := wstool.GetLogger(t.session)
-			logger.WithField("clientIP", t.session.Request.RemoteAddr).Info("user dropped! close connection!")
+			logger.WithField("clientIP", t.ipStr).Info("user dropped! close connection!")
 			t.session.Close()
 			t.Unlock()
 			t.Close()
@@ -259,15 +258,15 @@ func (t *Taos) waitSignal() {
 			}
 			whitelist, err := tool.GetWhitelist(t.conn)
 			if err != nil {
-				wstool.GetLogger(t.session).WithField("clientIP", t.session.Request.RemoteAddr).WithError(err).Errorln("get whitelist error! close connection!")
+				wstool.GetLogger(t.session).WithField("clientIP", t.ipStr).WithError(err).Errorln("get whitelist error! close connection!")
 				t.session.Close()
 				t.Unlock()
 				t.Close()
 				return
 			}
 			valid := tool.CheckWhitelist(whitelist, t.ip)
 			if !valid {
-				wstool.GetLogger(t.session).WithField("clientIP", t.session.Request.RemoteAddr).Errorln("ip not in whitelist! close connection!")
+				wstool.GetLogger(t.session).WithField("clientIP", t.ipStr).Errorln("ip not in whitelist! close connection!")
 				t.session.Close()
 				t.Unlock()
 				t.Close()

diff --git a/controller/ws/schemaless/schemaless.go b/controller/ws/schemaless/schemaless.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"net"
-	"strings"
 	"sync"
 	"time"
 	"unsafe"
@@ -21,6 +20,7 @@ import (
 	"github.com/taosdata/taosadapter/v3/db/tool"
 	"github.com/taosdata/taosadapter/v3/log"
 	"github.com/taosdata/taosadapter/v3/thread"
+	"github.com/taosdata/taosadapter/v3/tools/iptool"
 )
 
 type SchemalessController struct {
@@ -137,19 +137,20 @@ type TaosSchemaless struct {
 	dropUserNotify      chan struct{}
 	session             *melody.Session
 	ip                  net.IP
+	ipStr               string
 	wg                  sync.WaitGroup
 	sync.Mutex
 }
 
 func NewTaosSchemaless(session *melody.Session) *TaosSchemaless {
-	host, _, _ := net.SplitHostPort(strings.TrimSpace(session.Request.RemoteAddr))
-	ipAddr := net.ParseIP(host)
+	ipAddr := iptool.GetRealIP(session.Request)
 	return &TaosSchemaless{
 		exit:                make(chan struct{}),
 		whitelistChangeChan: make(chan int64, 1),
 		dropUserNotify:      make(chan struct{}, 1),
 		session:             session,
 		ip:                  ipAddr,
+		ipStr:               ipAddr.String(),
 	}
 }
 
@@ -166,7 +167,7 @@ func (t *TaosSchemaless) waitSignal() {
 				return
 			}
 			logger := wstool.GetLogger(t.session)
-			logger.WithField("clientIP", t.session.Request.RemoteAddr).Info("user dropped! close connection!")
+			logger.WithField("clientIP", t.ipStr).Info("user dropped! close connection!")
 			t.session.Close()
 			t.Unlock()
 			t.close()
@@ -179,14 +180,14 @@ func (t *TaosSchemaless) waitSignal() {
 			}
 			whitelist, err := tool.GetWhitelist(t.conn)
 			if err != nil {
-				wstool.GetLogger(t.session).WithField("clientIP", t.session.Request.RemoteAddr).WithError(err).Errorln("get whitelist error! close connection!")
+				wstool.GetLogger(t.session).WithField("clientIP", t.ipStr).WithError(err).Errorln("get whitelist error! close connection!")
 				t.session.Close()
 				t.Unlock()
 				return
 			}
 			valid := tool.CheckWhitelist(whitelist, t.ip)
 			if !valid {
-				wstool.GetLogger(t.session).WithField("clientIP", t.session.Request.RemoteAddr).Errorln("ip not in whitelist! close connection!")
+				wstool.GetLogger(t.session).WithField("clientIP", t.ipStr).Errorln("ip not in whitelist! close connection!")
 				t.session.Close()
 				t.Unlock()
 				t.close()

diff --git a/controller/ws/stmt/stmt.go b/controller/ws/stmt/stmt.go
@@ -7,7 +7,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"net"
-	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -30,6 +29,7 @@ import (
 	"github.com/taosdata/taosadapter/v3/log"
 	"github.com/taosdata/taosadapter/v3/thread"
 	"github.com/taosdata/taosadapter/v3/tools"
+	"github.com/taosdata/taosadapter/v3/tools/iptool"
 )
 
 type STMTController struct {
@@ -251,6 +251,7 @@ type TaosStmt struct {
 	dropUserNotify      chan struct{}
 	session             *melody.Session
 	ip                  net.IP
+	ipStr               string
 	wg                  sync.WaitGroup
 	sync.Mutex
 }
@@ -268,7 +269,7 @@ func (t *TaosStmt) waitSignal() {
 				return
 			}
 			logger := wstool.GetLogger(t.session)
-			logger.WithField("clientIP", t.session.Request.RemoteAddr).Info("user dropped! close connection!")
+			logger.WithField("clientIP", t.ipStr).Info("user dropped! close connection!")
 			t.session.Close()
 			t.Unlock()
 			t.Close()
@@ -281,15 +282,15 @@ func (t *TaosStmt) waitSignal() {
 			}
 			whitelist, err := tool.GetWhitelist(t.conn)
 			if err != nil {
-				wstool.GetLogger(t.session).WithField("clientIP", t.session.Request.RemoteAddr).WithError(err).Errorln("get whitelist error! close connection!")
+				wstool.GetLogger(t.session).WithField("clientIP", t.ipStr).WithError(err).Errorln("get whitelist error! close connection!")
 				t.session.Close()
 				t.Unlock()
 				t.Close()
 				return
 			}
 			valid := tool.CheckWhitelist(whitelist, t.ip)
 			if !valid {
-				wstool.GetLogger(t.session).WithField("clientIP", t.session.Request.RemoteAddr).Errorln("ip not in whitelist! close connection!")
+				wstool.GetLogger(t.session).WithField("clientIP", t.ipStr).Errorln("ip not in whitelist! close connection!")
 				t.session.Close()
 				t.Unlock()
 				t.Close()
@@ -319,15 +320,15 @@ func (s *StmtItem) clean() {
 }
 
 func NewTaosStmt(session *melody.Session) *TaosStmt {
-	host, _, _ := net.SplitHostPort(strings.TrimSpace(session.Request.RemoteAddr))
-	ipAddr := net.ParseIP(host)
+	ipAddr := iptool.GetRealIP(session.Request)
 	return &TaosStmt{
 		StmtList:            list.New(),
 		exit:                make(chan struct{}),
 		whitelistChangeChan: make(chan int64, 1),
 		dropUserNotify:      make(chan struct{}, 1),
 		session:             session,
 		ip:                  ipAddr,
+		ipStr:               ipAddr.String(),
 	}
 }
 

diff --git a/controller/ws/tmq/tmq.go b/controller/ws/tmq/tmq.go
@@ -8,7 +8,6 @@ import (
 	"errors"
 	"net"
 	"strconv"
-	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -29,6 +28,7 @@ import (
 	"github.com/taosdata/taosadapter/v3/db/tool"
 	"github.com/taosdata/taosadapter/v3/httperror"
 	"github.com/taosdata/taosadapter/v3/log"
+	"github.com/taosdata/taosadapter/v3/tools/iptool"
 	"github.com/taosdata/taosadapter/v3/tools/jsontype"
 )
 
@@ -244,6 +244,7 @@ type TMQ struct {
 	whitelistChangeChan chan int64
 	session             *melody.Session
 	ip                  net.IP
+	ipStr               string
 	wg                  sync.WaitGroup
 	conn                unsafe.Pointer
 	sync.Mutex
@@ -260,8 +261,7 @@ type Message struct {
 }
 
 func NewTaosTMQ(session *melody.Session) *TMQ {
-	host, _, _ := net.SplitHostPort(strings.TrimSpace(session.Request.RemoteAddr))
-	ipAddr := net.ParseIP(host)
+	ipAddr := iptool.GetRealIP(session.Request)
 	return &TMQ{
 		tmpMessage:          &Message{buffer: &bytes.Buffer{}},
 		handler:             tmqhandle.GlobalTMQHandlerPoll.Get(),
@@ -273,6 +273,7 @@ func NewTaosTMQ(session *melody.Session) *TMQ {
 		dropUserNotify:      make(chan struct{}, 1),
 		session:             session,
 		ip:                  ipAddr,
+		ipStr:               ipAddr.String(),
 	}
 }
 
@@ -289,7 +290,7 @@ func (t *TMQ) waitSignal() {
 				return
 			}
 			logger := wstool.GetLogger(t.session)
-			logger.WithField("clientIP", t.session.Request.RemoteAddr).Info("user dropped! close connection!")
+			logger.WithField("clientIP", t.ipStr).Info("user dropped! close connection!")
 			t.session.Close()
 			t.Unlock()
 			t.Close(logger)
@@ -303,7 +304,7 @@ func (t *TMQ) waitSignal() {
 			whitelist, err := tool.GetWhitelist(t.conn)
 			if err != nil {
 				logger := wstool.GetLogger(t.session)
-				logger.WithField("clientIP", t.session.Request.RemoteAddr).WithError(err).Errorln("get whitelist error! close connection!")
+				logger.WithField("clientIP", t.ipStr).WithError(err).Errorln("get whitelist error! close connection!")
 				t.session.Close()
 				t.Unlock()
 				t.Close(logger)
@@ -312,7 +313,7 @@ func (t *TMQ) waitSignal() {
 			valid := tool.CheckWhitelist(whitelist, t.ip)
 			if !valid {
 				logger := wstool.GetLogger(t.session)
-				logger.WithField("clientIP", t.session.Request.RemoteAddr).Errorln("ip not in whitelist! close connection!")
+				logger.WithField("clientIP", t.ipStr).Errorln("ip not in whitelist! close connection!")
 				t.session.Close()
 				t.Unlock()
 				t.Close(logger)

diff --git a/controller/ws/ws/handler.go b/controller/ws/ws/handler.go
@@ -8,7 +8,6 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"strings"
 	"sync"
 	"time"
 	"unsafe"
@@ -30,6 +29,7 @@ import (
 	"github.com/taosdata/taosadapter/v3/monitor"
 	"github.com/taosdata/taosadapter/v3/thread"
 	"github.com/taosdata/taosadapter/v3/tools"
+	"github.com/taosdata/taosadapter/v3/tools/iptool"
 	"github.com/taosdata/taosadapter/v3/tools/jsontype"
 	"github.com/taosdata/taosadapter/v3/version"
 )
@@ -53,8 +53,7 @@ type messageHandler struct {
 }
 
 func newHandler(session *melody.Session) *messageHandler {
-	host, _, _ := net.SplitHostPort(strings.TrimSpace(session.Request.RemoteAddr))
-	ipAddr := net.ParseIP(host)
+	ipAddr := iptool.GetRealIP(session.Request)
 	return &messageHandler{
 		queryResults:        NewQueryResultHolder(),
 		stmts:               NewStmtHolder(),
@@ -79,7 +78,7 @@ func (h *messageHandler) waitSignal() {
 				h.Unlock()
 				return
 			}
-			wstool.GetLogger(h.session).WithField("clientIP", h.session.Request.RemoteAddr).Info("user dropped! close connection!")
+			wstool.GetLogger(h.session).WithField("clientIP", h.ipStr).Info("user dropped! close connection!")
 			h.session.Close()
 			h.Unlock()
 			h.Close()
@@ -92,15 +91,15 @@ func (h *messageHandler) waitSignal() {
 			}
 			whitelist, err := tool.GetWhitelist(h.conn)
 			if err != nil {
-				wstool.GetLogger(h.session).WithField("clientIP", h.session.Request.RemoteAddr).WithError(err).Errorln("get whitelist error! close connection!")
+				wstool.GetLogger(h.session).WithField("clientIP", h.ipStr).WithError(err).Errorln("get whitelist error! close connection!")
 				h.session.Close()
 				h.Unlock()
 				h.Close()
 				return
 			}
 			valid := tool.CheckWhitelist(whitelist, h.ip)
 			if !valid {
-				wstool.GetLogger(h.session).WithField("clientIP", h.session.Request.RemoteAddr).Errorln("ip not in whitelist! close connection!")
+				wstool.GetLogger(h.session).WithField("clientIP", h.ipStr).Errorln("ip not in whitelist! close connection!")
 				h.session.Close()
 				h.Unlock()
 				h.Close()