fix: ClientIP() handles uds instead of returning empty

pkg/app/context.go

@@ -99,43 +99,43 @@ var defaultClientIPOptions = ClientIPOptions{
 	TrustedCIDRs:    defaultTrustedCIDRs,
 }
 
+var loopbackIP = net.ParseIP("127.0.0.1")
+
 // ClientIPWithOption used to generate custom ClientIP function and set by engine.SetClientIPFunc
 func ClientIPWithOption(opts ClientIPOptions) ClientIP {
 	return func(ctx *RequestContext) string {
-		RemoteIPHeaders := opts.RemoteIPHeaders
-		TrustedCIDRs := opts.TrustedCIDRs
-
-		remoteIPStr, _, err := net.SplitHostPort(strings.TrimSpace(ctx.RemoteAddr().String()))
-		if err != nil {
-			return ""
-		}
-
-		remoteIP := net.ParseIP(remoteIPStr)
-		if remoteIP == nil {
-			return ""
+		remoteIPStr := ""
+		trustedProxy := false
+		if addr := ctx.RemoteAddr(); strings.HasPrefix(addr.Network(), "unix") {
+			// unix, unixgram, unixpacket is considered same as "127.0.0.1"
+			remoteIPStr = addr.String()
+			trustedProxy = isTrustedProxy(opts.TrustedCIDRs, loopbackIP)
+		} else {
+			h, _, err := net.SplitHostPort(strings.TrimSpace(addr.String()))
+			if err != nil {
+				return ""
+			}
+			remoteIPStr = h
+			trustedProxy = isTrustedProxy(opts.TrustedCIDRs, net.ParseIP(h))
 		}
 
-		trusted := isTrustedProxy(TrustedCIDRs, remoteIP)
-
-		if trusted {
-			for _, headerName := range RemoteIPHeaders {
-				ip, valid := validateHeader(TrustedCIDRs, ctx.Request.Header.Get(headerName))
+		if trustedProxy {
+			for _, headerName := range opts.RemoteIPHeaders {
+				ip, valid := validateHeader(opts.TrustedCIDRs, ctx.Request.Header.Get(headerName))
 				if valid {
 					return ip
 				}
 			}
 		}
-
 		return remoteIPStr
 	}
 }
 
 // isTrustedProxy will check whether the IP address is included in the trusted list according to trustedCIDRs
 func isTrustedProxy(trustedCIDRs []*net.IPNet, remoteIP net.IP) bool {
-	if trustedCIDRs == nil {
+	if trustedCIDRs == nil || remoteIP == nil {
 		return false
 	}
-
 	for _, cidr := range trustedCIDRs {
 		if cidr.Contains(remoteIP) {
 			return true

pkg/app/context_test.go

@@ -848,32 +848,27 @@ func TestContextContentType(t *testing.T) {
 	assert.DeepEqual(t, consts.MIMEApplicationJSONUTF8, bytesconv.B2s(c.ContentType()))
 }
 
-type MockIpConn struct {
+type MockConn struct {
 	*mock.Conn
-	RemoteIp string
-	Port     int
+
+	remote net.Addr
 }
 
-func (c *MockIpConn) RemoteAddr() net.Addr {
-	return &net.UDPAddr{
-		IP:   net.ParseIP(c.RemoteIp),
-		Port: c.Port,
-	}
+func (c *MockConn) RemoteAddr() net.Addr {
+	return c.remote
 }
 
 func newContextClientIPTest() *RequestContext {
 	c := NewContext(0)
-	c.conn = &MockIpConn{
-		Conn:     mock.NewConn(""),
-		RemoteIp: "127.0.0.1",
-		Port:     8080,
+	c.conn = &MockConn{
+		remote: &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 8080},
 	}
 	c.Request.Header.Set("X-Real-IP", " 10.10.10.10  ")
 	c.Request.Header.Set("X-Forwarded-For", "  20.20.20.20, 30.30.30.30")
 	return c
 }
 
-func TestClientIp(t *testing.T) {
+func TestClientIP(t *testing.T) {
 	c := newContextClientIPTest()
 	// default X-Forwarded-For and X-Real-IP behaviour
 	assert.DeepEqual(t, "20.20.20.20", c.ClientIP())
@@ -908,6 +903,14 @@ func TestClientIp(t *testing.T) {
 	}
 	c.SetClientIPFunc(ClientIPWithOption(opts))
 	assert.DeepEqual(t, "30.30.30.30", c.ClientIP())
+
+	// UDS
+	c.conn = &MockConn{remote: &net.UnixAddr{Net: "unix", Name: "/tmp/test.sock"}}
+	assert.DeepEqual(t, "30.30.30.30", c.ClientIP())
+
+	// err: Addr not host:port
+	c.conn = &MockConn{remote: &net.UnixAddr{Net: "tcp", Name: "/tmp/test.sock"}}
+	assert.DeepEqual(t, "", c.ClientIP())
 }
 
 func TestSetClientIPFunc(t *testing.T) {