-
Notifications
You must be signed in to change notification settings - Fork 2
/
client.go
127 lines (120 loc) · 2.96 KB
/
client.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
package client
import (
"context"
"fmt"
"log"
"net"
"strconv"
"time"
"github.com/sirupsen/logrus"
"github.com/spf13/cobra"
"github.com/winglq/l4proxy/src/api"
"github.com/winglq/l4proxy/src/handler"
"github.com/winglq/l4proxy/src/port_map"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
)
type Options struct {
SvrAddr string
PubPort int32
IntPort int32
Name string
SharePub bool
BackendPort int32
}
func createClient(client api.ControlServiceClient, opt *Options, backendPort int32) (api.ControlService_CreateClientClient, error) {
for {
c, err := client.CreateClient(context.TODO(), &api.CreateClientRequest{
DisplayName: opt.Name,
PublicPort: opt.PubPort,
InternalPort: opt.IntPort,
SharePublicAddr: opt.SharePub,
Protocol: "tcp",
BackendPort: backendPort,
})
if err == nil {
return c, nil
}
if grpc.Code(err) == codes.Unavailable {
log.Printf("reconnecting after 5 second due to err: %v", err)
time.Sleep(5 * time.Second)
continue
} else if grpc.Code(err) == codes.Canceled {
return nil, err
} else if err != nil {
panic(err)
}
}
}
func CreateRunFunc(done chan struct{}, opt *Options, onNewConn func(resp *api.Client, host, port string) (*handler.PairedConn, error)) func(cmd *cobra.Command, args []string) {
ret := func(cmd *cobra.Command, args []string) {
port := "22"
host := "127.0.0.1"
if len(args) > 0 {
host = args[0]
}
if len(args) > 1 {
port = args[1]
}
c, err := grpc.Dial(opt.SvrAddr, grpc.WithInsecure())
if err != nil {
panic(err)
}
go func() {
<-done
c.Close()
}()
client := api.NewControlServiceClient(c)
backendPort := opt.BackendPort
var pt int64
if backendPort == 0 {
pt, err = strconv.ParseInt(port, 10, 32)
if err != nil {
log.Fatalf("backend port format error")
}
backendPort = int32(pt)
}
mapper := port_map.NewDummyPortMapper()
mapper.MapPort("", int32(pt), "tcp", backendPort)
conClient, err := createClient(client, opt, backendPort)
if err != nil && grpc.Code(err) == codes.Canceled {
return
} else if err != nil {
panic(err)
}
for {
resp, err := conClient.Recv()
if err != nil {
if grpc.Code(err) == codes.Canceled {
break
} else {
logrus.Errorf("recv messsage failed: %v", err)
conClient, err = createClient(client, opt, backendPort)
if err != nil && grpc.Code(err) == codes.Canceled {
mapper.UnmapPort("tcp", backendPort)
return
}
continue
}
}
if resp.InternalAddress != "" {
pair, err := onNewConn(resp, host, port)
if err != nil {
fmt.Printf("create pair failed: %v\n", err)
}
if pair != nil {
defer pair.Close()
}
} else {
fmt.Printf("PUBLIC ADDRESS: %s\n", resp.PublicAddress)
_, port, err := net.SplitHostPort(resp.PublicAddress)
if err != nil {
panic(err)
}
p, _ := strconv.ParseInt(port, 10, 64)
opt.PubPort = int32(p)
}
}
}
return ret
}