按协议判断连接权限，优化权限检查效率

2025-03-08 11:40:52 +08:00
parent 5786ac9d99
commit f996a20823
11 changed files with 328 additions and 101 deletions
--- a/server/fwd/auth/auth_test.go
+++ b/server/fwd/auth/auth_test.go
@@ -0,0 +1,151 @@
+package auth
+
+import (
+	"net"
+	"proxy-server/server/fwd/core"
+	"proxy-server/server/pkg/orm"
+	"reflect"
+	"testing"
+	"time"
+)
+
+type MockConn struct {
+	local  net.Addr
+	remote net.Addr
+}
+
+func (m MockConn) LocalAddr() net.Addr {
+	return m.local
+}
+
+func (m MockConn) RemoteAddr() net.Addr {
+	return m.remote
+}
+
+func (m MockConn) Read(b []byte) (n int, err error) {
+	return 0, nil
+}
+
+func (m MockConn) Write(b []byte) (n int, err error) {
+	return 0, nil
+}
+
+func (m MockConn) Close() error {
+	return nil
+}
+
+func (m MockConn) SetDeadline(t time.Time) error {
+	return nil
+}
+
+func (m MockConn) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (m MockConn) SetWriteDeadline(t time.Time) error {
+	return nil
+}
+
+func TestCheckIp(t *testing.T) {
+
+	orm.InitForTest(
+		"host=localhost port=5432 user=proxy password=proxy dbname=app sslmode=disable TimeZone=Asia/Shanghai",
+	)
+
+	type args struct {
+		conn  net.Conn
+		proto Protocol
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    *core.AuthContext
+		wantErr bool
+	}{
+		{
+			name: "test-ok",
+			args: args{
+				conn: MockConn{
+					remote: &net.TCPAddr{
+						IP:   []byte{127, 0, 0, 1},
+						Port: 12345,
+					},
+					local: &net.TCPAddr{
+						IP:   []byte{127, 0, 0, 1},
+						Port: 20001,
+					},
+				},
+				proto: Http,
+			},
+			want: &core.AuthContext{
+				Payload: core.Payload{ID: 1},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := CheckIp(tt.args.conn, tt.args.proto)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("CheckIp() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got.Payload, tt.want.Payload) || !reflect.DeepEqual(got.Meta, tt.want.Meta) {
+				t.Errorf("CheckIp() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestCheckPass(t *testing.T) {
+
+	orm.InitForTest(
+		"host=localhost port=5432 user=proxy password=proxy dbname=app sslmode=disable TimeZone=Asia/Shanghai",
+	)
+
+	type args struct {
+		conn     net.Conn
+		username string
+		password string
+		proto    Protocol
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    *core.AuthContext
+		wantErr bool
+	}{
+		{
+			name: "test-ok",
+			args: args{
+				conn: MockConn{
+					remote: &net.TCPAddr{
+						IP:   []byte{127, 0, 0, 1},
+						Port: 12345,
+					},
+					local: &net.TCPAddr{
+						IP:   []byte{127, 0, 0, 1},
+						Port: 20001,
+					},
+				},
+				proto:    Http,
+				username: "ip4http",
+				password: "asdf",
+			},
+			want: &core.AuthContext{
+				Payload: core.Payload{ID: 1},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := CheckPass(tt.args.conn, tt.args.proto, tt.args.username, tt.args.password)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("CheckPass() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got.Payload, tt.want.Payload) || !reflect.DeepEqual(got.Meta, tt.want.Meta) {
+				t.Errorf("CheckPass() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}