proxy/server/fwd/auth/auth_test.go

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)
			}
		})
	}
}