-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
af11f97
commit 96ae108
Showing
5 changed files
with
247 additions
and
57 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| package internal | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "net" | ||
| "sync" | ||
| ) | ||
|
|
||
| type ConnectionStore struct { | ||
| conns map[string]*WsConnAdapter | ||
| mu sync.RWMutex | ||
| } | ||
|
|
||
| func NewConnectionStore() *ConnectionStore { | ||
| return &ConnectionStore{ | ||
| conns: make(map[string]*WsConnAdapter), | ||
| } | ||
| } | ||
|
|
||
| // ConnCacheSet store connection in cache. | ||
| func (cs *ConnectionStore) ConnCacheSet(addr net.Addr, conn *WsConnAdapter) error { | ||
| h, err := extractIpFromAddr(addr) | ||
| if err != nil { | ||
| return fmt.Errorf("cant set conn cache for %s, %w", addr.String(), err) | ||
| } | ||
|
|
||
| cs.mu.Lock() | ||
| defer cs.mu.Unlock() | ||
|
|
||
| cs.conns[h] = conn | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| // ConnCacheRemove remove connection from cache. | ||
| func (cs *ConnectionStore) ConnCacheRemove(addr net.Addr) bool { | ||
| h, err := extractIpFromAddr(addr) | ||
| if err != nil { | ||
| return false | ||
| } | ||
|
|
||
| cs.mu.Lock() | ||
| defer cs.mu.Unlock() | ||
|
|
||
| delete(cs.conns, h) | ||
|
|
||
| return true | ||
| } | ||
|
|
||
| // ConnCacheGet get connection from cache. | ||
| func (cs *ConnectionStore) ConnCacheGet(addr net.Addr) (*WsConnAdapter, bool, error) { | ||
| h, err := extractIpFromAddr(addr) | ||
| if err != nil { | ||
| return nil, false, fmt.Errorf("cant get conn for addr %s from cache, %w", addr.String(), err) | ||
| } | ||
|
|
||
| cs.mu.Lock() | ||
| defer cs.mu.Unlock() | ||
|
|
||
| conn, ok := cs.conns[h] | ||
| if !ok { | ||
| return nil, false, nil | ||
| } | ||
|
|
||
| return conn, true, nil | ||
| } | ||
|
|
||
| func extractIpFromAddr(addr net.Addr) (string, error) { | ||
| h, _, err := net.SplitHostPort(addr.String()) | ||
| if err != nil { | ||
| return "", err | ||
| } | ||
|
|
||
| return h, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,164 @@ | ||
| package internal | ||
|
|
||
| import ( | ||
| "errors" | ||
| "net" | ||
| "sync" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/assert" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| func TestConnectionStore_ConnCacheSet(t *testing.T) { | ||
| type fields struct { | ||
| conns map[string]*WsConnAdapter | ||
| mu sync.RWMutex | ||
| } | ||
| type args struct { | ||
| addr net.Addr | ||
| conn *WsConnAdapter | ||
| } | ||
| tests := []struct { | ||
| name string | ||
| fields fields | ||
| expectedAddr string | ||
| args args | ||
| err error | ||
| }{ | ||
| { | ||
| name: "Set ok", | ||
| fields: fields{ | ||
| conns: make(map[string]*WsConnAdapter), | ||
| }, | ||
| args: args{ | ||
| addr: &WsAddr{ | ||
| WsAddrTCP: func() net.TCPAddr { | ||
| a, _ := net.ResolveTCPAddr("tcp", "192.168.1.1:8888") | ||
| return *a | ||
| }(), | ||
| }, | ||
| conn: &WsConnAdapter{ | ||
| realRemoteAddr: func() *net.TCPAddr { | ||
| a, _ := net.ResolveTCPAddr("tcp", "192.168.1.1:8888") | ||
| return a | ||
| }(), | ||
| }, | ||
| }, | ||
| expectedAddr: "192.168.1.1:8888", | ||
| err: nil, | ||
| }, | ||
| { | ||
| name: "Set cant parse addr", | ||
| fields: fields{ | ||
| conns: make(map[string]*WsConnAdapter), | ||
| }, | ||
| args: args{ | ||
| addr: &WsAddr{ | ||
| WsAddrTCP: func() net.TCPAddr { | ||
| a := net.TCPAddr{ | ||
| IP: net.IPv4(192, 168, 1, 1), | ||
| } | ||
| return a | ||
| }(), | ||
| }, | ||
| conn: &WsConnAdapter{ | ||
| realRemoteAddr: func() *net.TCPAddr { | ||
| a := net.TCPAddr{ | ||
| IP: net.IPv4(192, 168, 1, 1), | ||
| } | ||
| return &a | ||
| }(), | ||
| }, | ||
| }, | ||
| expectedAddr: "192.168.1.1:0", | ||
| err: errors.New(""), | ||
| }, | ||
| } | ||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| cs := &ConnectionStore{ | ||
| conns: tt.fields.conns, | ||
| mu: tt.fields.mu, | ||
| } | ||
| err := cs.ConnCacheSet(tt.args.addr, tt.args.conn) | ||
| if err == nil { | ||
| require.NoError(t, err) | ||
| } else { | ||
| require.Errorf(t, err, err.Error()) | ||
| } | ||
| require.NotNil(t, tt.args.conn) | ||
| assert.Equal(t, tt.expectedAddr, tt.args.conn.RemoteAddr().String()) | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| func TestConnectionStore_ConnCacheGet(t *testing.T) { | ||
| type fields struct { | ||
| conns map[string]*WsConnAdapter | ||
| mu sync.RWMutex | ||
| } | ||
| type args struct { | ||
| addr net.Addr | ||
| conn *WsConnAdapter | ||
| setConnectionToCache bool | ||
| } | ||
| tests := []struct { | ||
| name string | ||
| fields fields | ||
| args args | ||
| want *WsConnAdapter | ||
| want1 bool | ||
| wantErr bool | ||
| }{ | ||
| { | ||
| name: "Get connection ok", | ||
| fields: fields{ | ||
| conns: make(map[string]*WsConnAdapter), | ||
| }, | ||
| args: args{ | ||
| addr: &WsAddr{ | ||
| WsAddrTCP: net.TCPAddr{ | ||
| IP: net.IPv4(192, 168, 1, 1), | ||
| }, | ||
| }, | ||
| setConnectionToCache: true, | ||
| conn: &WsConnAdapter{}, | ||
| }, | ||
| want: &WsConnAdapter{}, | ||
| want1: true, | ||
| wantErr: false, | ||
| }, | ||
| { | ||
| name: "Get connection when connection not set", | ||
| fields: fields{ | ||
| conns: make(map[string]*WsConnAdapter), | ||
| }, | ||
| args: args{ | ||
| addr: &WsAddr{ | ||
| WsAddrTCP: net.TCPAddr{ | ||
| IP: net.IPv4(192, 168, 1, 1), | ||
| }, | ||
| }, | ||
| setConnectionToCache: false, | ||
| conn: &WsConnAdapter{}, | ||
| }, | ||
| want: &WsConnAdapter{}, | ||
| want1: false, | ||
| wantErr: false, | ||
| }, | ||
| } | ||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| cs := &ConnectionStore{ | ||
| conns: tt.fields.conns, | ||
| } | ||
| if tt.args.setConnectionToCache { | ||
| cs.ConnCacheSet(tt.args.addr, tt.want) | ||
| } | ||
| _, ok, err := cs.ConnCacheGet(tt.args.addr) | ||
| require.Equal(t, tt.wantErr, err != nil) | ||
| assert.Equal(t, tt.want1, ok) | ||
| }) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters