Skip to content

Commit a861e27

Browse files
LeeEircLeeEirc
committed
feat: 增加 ssh forward
1 parent 0b07aff commit a861e27

16 files changed

Lines changed: 976 additions & 195 deletions

File tree

‎Makefile‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,6 @@ proto: proto-go
7474
.PHONY: proto-go
7575
proto-go:
7676
@mkdir -p ${proto_go_dir}/protobuf
77-
rm -rf ${proto_go_dir}/protobuf/*
7877
protoc --proto_path=${proto_path} --go_out=${proto_go_dir} \
7978
--go-grpc_out=${proto_go_dir} \
8079
--go_opt=paths=import \
@@ -85,7 +84,6 @@ proto-go:
8584
.PHONY: proto-py
8685
proto-py:
8786
@mkdir -p ./protobuf-py/protobuf
88-
rm -rf ./protobuf-py/protobuf/*
8987
python -m grpc_tools.protoc --proto_path=${proto_path} --python_out=./protobuf-py/protobuf \
9088
--grpc_python_out=./protobuf-py/protobuf \
9189
${proto_files}

‎cmd/common/forward_service.go‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package common
2+
3+
import (
4+
"errors"
5+
6+
"sync"
7+
8+
"golang.org/x/crypto/ssh"
9+
10+
"github.com/jumpserver/wisp/pkg/forward"
11+
"github.com/jumpserver/wisp/pkg/jms-sdk-go/model"
12+
"github.com/jumpserver/wisp/pkg/logger"
13+
"github.com/jumpserver/wisp/pkg/sshclient"
14+
15+
pb "github.com/jumpserver/wisp/protobuf-go/protobuf"
16+
)
17+
18+
func NewForwardCache() *ForwardCache {
19+
return &ForwardCache{
20+
lns: make(map[string]*forward.SSHForward),
21+
}
22+
}
23+
24+
type ForwardCache struct {
25+
lns map[string]*forward.SSHForward
26+
27+
sync.Mutex
28+
}
29+
30+
func (s *ForwardCache) Add(key string, forward *forward.SSHForward) {
31+
s.Lock()
32+
defer s.Unlock()
33+
s.lns[key] = forward
34+
35+
}
36+
37+
func (s *ForwardCache) Remove(key string) {
38+
s.Lock()
39+
defer s.Unlock()
40+
delete(s.lns, key)
41+
42+
}
43+
44+
func (s *ForwardCache) Get(key string) *forward.SSHForward {
45+
s.Lock()
46+
defer s.Unlock()
47+
return s.lns[key]
48+
}
49+
50+
var (
51+
ErrNoAvailable = errors.New("no available gateway")
52+
)
53+
54+
func FindAvailableDomainGateway(domain *model.Domain) (*ssh.Client, error) {
55+
for i := range domain.Gateways {
56+
gateway := domain.Gateways[i]
57+
opts := make([]sshclient.Option, 0, 7)
58+
opts = append(opts, sshclient.WithHost(gateway.IP))
59+
opts = append(opts, sshclient.WithPort(gateway.Port))
60+
opts = append(opts, sshclient.WithUsername(gateway.Username))
61+
opts = append(opts, sshclient.WithPassword(gateway.Password))
62+
opts = append(opts, sshclient.WithPrivateKey(gateway.PrivateKey))
63+
opts = append(opts, sshclient.WithPassphrase(gateway.Password))
64+
opts = append(opts, sshclient.WithTimeout(15))
65+
proxyClient, err := sshclient.New(opts...)
66+
if err == nil {
67+
return proxyClient, nil
68+
}
69+
logger.Infof("Domain %s use gateway %s failed: %s",
70+
domain.Name, gateway.Name, err)
71+
}
72+
logger.Errorf("Domain %s find available gateway failed: %s", domain.Name)
73+
return nil, ErrNoAvailable
74+
}
75+
76+
func FindAvailableGateway(gateways []*pb.Gateway) (*ssh.Client, error) {
77+
for i := range gateways {
78+
gateway := gateways[i]
79+
opts := make([]sshclient.Option, 0, 7)
80+
opts = append(opts, sshclient.WithHost(gateway.Ip))
81+
opts = append(opts, sshclient.WithPort(int(gateway.Port)))
82+
opts = append(opts, sshclient.WithUsername(gateway.Username))
83+
opts = append(opts, sshclient.WithPassword(gateway.Password))
84+
opts = append(opts, sshclient.WithPrivateKey(gateway.PrivateKey))
85+
opts = append(opts, sshclient.WithPassphrase(gateway.Password))
86+
opts = append(opts, sshclient.WithTimeout(15))
87+
proxyClient, err := sshclient.New(opts...)
88+
if err == nil {
89+
return proxyClient, nil
90+
}
91+
}
92+
return nil, ErrNoAvailable
93+
}

‎cmd/common/uploader_service.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ type UploaderService struct {
3434
commandCfg atomic.Value // model.CommandConfig
3535
replayCfg atomic.Value // model.ReplayConfig
3636

37+
terminalCfg atomic.Value // *model.TerminalConfig
38+
3739
wg sync.WaitGroup
3840
}
3941

@@ -66,6 +68,7 @@ func (u *UploaderService) watchConfig() {
6668
func (u *UploaderService) updateBackendCfg(termCfg *model.TerminalConfig) {
6769
u.commandCfg.Store(termCfg.CommandStorage)
6870
u.replayCfg.Store(termCfg.ReplayStorage)
71+
u.terminalCfg.Store(termCfg)
6972
}
7073

7174
func (u *UploaderService) getCommandBackend() CommandStorage {
@@ -78,6 +81,11 @@ func (u *UploaderService) getReplayBackend() ReplayStorage {
7881
return NewReplayBackend(u.apiClient, &cfg)
7982
}
8083

84+
func (u *UploaderService) GetTerminalSetting() model.TerminalConfig {
85+
cfg := u.terminalCfg.Load().(*model.TerminalConfig)
86+
return *cfg
87+
}
88+
8189
func (u *UploaderService) run() {
8290
cmdList := make([]*model.Command, 0, 10)
8391
notificationList := make([]*model.Command, 0, 10)

‎cmd/impl/convert_pb.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,3 +180,7 @@ var pbTicketMap = map[string]pb.TicketState_State{
180180
model.TicketRejected: pb.TicketState_Rejected,
181181
model.TicketClosed: pb.TicketState_Closed,
182182
}
183+
184+
func ConvertToPbSetting(setting *model.TerminalConfig) *pb.ComponentSetting {
185+
return &pb.ComponentSetting{MaxIdleTime: int32(setting.MaxIdleTime)}
186+
}

‎cmd/impl/jms.go‎

Lines changed: 71 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,12 @@ import (
44
"context"
55
"fmt"
66
"io"
7+
"net"
8+
"strconv"
79

810
"github.com/jumpserver/wisp/cmd/common"
11+
"github.com/jumpserver/wisp/pkg/forward"
12+
modelCommon "github.com/jumpserver/wisp/pkg/jms-sdk-go/common"
913
"github.com/jumpserver/wisp/pkg/jms-sdk-go/model"
1014
"github.com/jumpserver/wisp/pkg/jms-sdk-go/service"
1115
"github.com/jumpserver/wisp/pkg/logger"
@@ -15,9 +19,10 @@ import (
1519
func NewJMServer(apiClient *service.JMService, uploader *common.UploaderService,
1620
beat *common.BeatService) *JMServer {
1721
return &JMServer{
18-
apiClient: apiClient,
19-
uploader: uploader,
20-
beat: beat,
22+
apiClient: apiClient,
23+
uploader: uploader,
24+
beat: beat,
25+
forwardStore: common.NewForwardCache(),
2126
}
2227
}
2328

@@ -27,6 +32,8 @@ type JMServer struct {
2732

2833
uploader *common.UploaderService
2934
beat *common.BeatService
35+
36+
forwardStore *common.ForwardCache
3037
}
3138

3239
func (j *JMServer) GetDBTokenAuthInfo(ctx context.Context, req *pb.DBTokenRequest) (*pb.DBTokenResponse, error) {
@@ -45,6 +52,7 @@ func (j *JMServer) GetDBTokenAuthInfo(ctx context.Context, req *pb.DBTokenReques
4552
logger.Error(msg)
4653
return &pb.DBTokenResponse{Status: &status}, nil
4754
}
55+
setting := j.uploader.GetTerminalSetting()
4856
dbTokenInfo := pb.DBTokenAuthInfo{
4957
KeyId: tokenAuthInfo.Id,
5058
SecreteId: tokenAuthInfo.Secret,
@@ -54,7 +62,8 @@ func (j *JMServer) GetDBTokenAuthInfo(ctx context.Context, req *pb.DBTokenReques
5462
SystemUser: ConvertToProtobufSystemUser(tokenAuthInfo.SystemUserAuthInfo),
5563
Permission: ConvertToProtobufPermission(model.Permission{Actions: tokenAuthInfo.Actions}),
5664
ExpireInfo: ConvertToProtobufExpireInfo(model.ExpireInfo{ExpireAt: tokenAuthInfo.ExpiredAt}),
57-
Gateways: ConvertToProtobufGateWays([]model.Gateway{tokenAuthInfo.Gateway}),
65+
Gateways: ConvertToProtobufGateWays(tokenAuthInfo.Domain.Gateways),
66+
Setting: ConvertToPbSetting(&setting),
5867
}
5968
status.Ok = true
6069
logger.Debugf("Get database auth info success by token: %s", req.Token)
@@ -260,3 +269,61 @@ func (j *JMServer) CheckOrCreateAssetLoginTicket(ctx context.Context,
260269
TicketInfo: ConvertToPbTicketInfo(&res.TicketInfo),
261270
Status: &status}, nil
262271
}
272+
273+
func (j *JMServer) CreateForward(ctx context.Context, req *pb.ForwardRequest) (*pb.ForwardResponse, error) {
274+
var (
275+
status pb.Status
276+
)
277+
host := req.GetHost()
278+
port := strconv.FormatInt(int64(req.GetPort()), 10)
279+
dstAddr := net.JoinHostPort(host, port)
280+
gateways := req.GetGateways()
281+
client, err := common.FindAvailableGateway(gateways)
282+
if err != nil {
283+
status.Err = err.Error()
284+
return &pb.ForwardResponse{
285+
Status: &status,
286+
}, nil
287+
}
288+
forwardProxy := forward.SSHForward{
289+
Client: client,
290+
DstAddr: dstAddr,
291+
}
292+
if err = forwardProxy.Start(); err != nil {
293+
status.Err = err.Error()
294+
_ = client.Close()
295+
logger.Errorf("Start forward proxy failed: %s", err)
296+
return &pb.ForwardResponse{
297+
Status: &status,
298+
}, nil
299+
}
300+
id := modelCommon.UUID()
301+
j.forwardStore.Add(id, &forwardProxy)
302+
lnAddr := forwardProxy.GetTCPAddr()
303+
status.Ok = true
304+
logger.Infof("Start forward proxy: id %s on %s", id, lnAddr.String())
305+
ret := &pb.ForwardResponse{
306+
Status: &status,
307+
Id: id,
308+
Host: lnAddr.IP.String(),
309+
Port: int32(lnAddr.Port),
310+
}
311+
return ret, nil
312+
}
313+
314+
func (j *JMServer) DeleteForward(ctx context.Context, req *pb.ForwardDeleteRequest) (*pb.StatusResponse, error) {
315+
var (
316+
status pb.Status
317+
)
318+
id := req.GetId()
319+
if forwardProxy := j.forwardStore.Get(id); forwardProxy != nil {
320+
forwardProxy.Stop()
321+
status.Ok = true
322+
j.forwardStore.Remove(id)
323+
logger.Infof("Forward remove id %s", id)
324+
return &pb.StatusResponse{Status: &status}, nil
325+
}
326+
status.Err = fmt.Sprintf("not found forward %s", id)
327+
logger.Errorf("Forward not found id %s", id)
328+
return &pb.StatusResponse{Status: &status}, nil
329+
}

‎pkg/forward/ssh.go‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package forward
2+
3+
import (
4+
"io"
5+
"net"
6+
7+
"golang.org/x/crypto/ssh"
8+
9+
"github.com/jumpserver/wisp/pkg/logger"
10+
)
11+
12+
type SSHForward struct {
13+
Client *ssh.Client
14+
DstAddr string
15+
16+
ln net.Listener
17+
addr *net.TCPAddr
18+
}
19+
20+
func (s *SSHForward) Start() error {
21+
ln, err := net.Listen("tcp", "0.0.0.0:0")
22+
if err != nil {
23+
return err
24+
}
25+
s.ln = ln
26+
go s.run()
27+
return nil
28+
}
29+
30+
func (s *SSHForward) GetTCPAddr() *net.TCPAddr {
31+
return s.ln.Addr().(*net.TCPAddr)
32+
}
33+
34+
func (s *SSHForward) Stop() {
35+
if s.ln != nil {
36+
if err := s.ln.Close(); err != nil {
37+
logger.Error(err)
38+
}
39+
}
40+
if err := s.Client.Close(); err != nil {
41+
logger.Error(err)
42+
}
43+
}
44+
45+
func (s *SSHForward) String() string {
46+
return s.DstAddr
47+
}
48+
49+
func (s *SSHForward) run() {
50+
for {
51+
conn, err := s.ln.Accept()
52+
if err != nil {
53+
if err != net.ErrClosed {
54+
logger.Errorf("listen %s accept failed: %v", s.ln.Addr(), err)
55+
}
56+
return
57+
}
58+
go s.forward(conn)
59+
}
60+
}
61+
62+
func (s *SSHForward) forward(conn net.Conn) {
63+
defer conn.Close()
64+
proxyCon, err := s.Client.Dial("tcp", s.DstAddr)
65+
if err != nil {
66+
logger.Errorf("ssh.Dial %s failed: %s\n", s.DstAddr, err)
67+
return
68+
}
69+
go func() {
70+
defer proxyCon.Close()
71+
if _, err = io.Copy(proxyCon, conn); err != nil && err != io.EOF {
72+
logger.Errorf("io.Copy local-> proxy err: %s\n", err)
73+
}
74+
}()
75+
if _, err = io.Copy(conn, proxyCon); err != nil && err != io.EOF {
76+
logger.Errorf("io.Copy proxy -> local err: %s\n", err)
77+
}
78+
}

‎pkg/jms-sdk-go/common/uuid.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,8 @@ func IsUUID(sid string) bool {
66
_, err := uuid.FromString(sid)
77
return err == nil
88
}
9+
10+
func UUID() string {
11+
ret, _ := uuid.NewV4()
12+
return ret.String()
13+
}

‎pkg/jms-sdk-go/service/jms_token.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ type TokenAuthInfo struct {
4848
Asset *model.Asset `json:"asset,omitempty"`
4949
ExpiredAt int64 `json:"expired_at"`
5050
Gateway model.Gateway `json:"gateway,omitempty"`
51+
Domain model.Domain `json:"domain"`
5152

5253
CmdFilterRules []model.FilterRule `json:"cmd_filter_rules,omitempty"`
5354

0 commit comments

Comments
 (0)