mirror of
https://github.com/yunionio/cloudpods.git
synced 2026-10-24 23:41:52 +08:00
cloudproxy: initial version
This commit is contained in:
@@ -0,0 +1,345 @@
|
||||
// Copyright 2019 Yunion
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package ssh
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/crypto/ssh"
|
||||
|
||||
"yunion.io/x/log"
|
||||
"yunion.io/x/pkg/errors"
|
||||
"yunion.io/x/pkg/util/sets"
|
||||
)
|
||||
|
||||
type addrMap map[string]interface{}
|
||||
type portMap map[int]addrMap
|
||||
|
||||
func (pm portMap) contains(port int, addr string) bool {
|
||||
am, ok := pm[port]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return am.contains(addr)
|
||||
}
|
||||
|
||||
func (pm portMap) get(port int, addr string) interface{} {
|
||||
am, ok := pm[port]
|
||||
if ok {
|
||||
return am.get(addr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (pm portMap) set(port int, addr string, v interface{}) {
|
||||
am, ok := pm[port]
|
||||
if !ok {
|
||||
am = addrMap{}
|
||||
pm[port] = am
|
||||
}
|
||||
am.set(addr, v)
|
||||
}
|
||||
|
||||
func (pm portMap) delete(port int, addr string) {
|
||||
if am, ok := pm[port]; ok {
|
||||
am.delete(addr)
|
||||
}
|
||||
}
|
||||
|
||||
func (am addrMap) contains(addr string) bool {
|
||||
const (
|
||||
ip4wild = "0.0.0.0"
|
||||
ip6wild = "::"
|
||||
)
|
||||
_, ok := am[addr]
|
||||
if ok {
|
||||
return true
|
||||
}
|
||||
if _, ok := am[ip4wild]; ok {
|
||||
return true
|
||||
}
|
||||
if _, ok := am[ip6wild]; ok {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (am addrMap) get(addr string) interface{} {
|
||||
return am[addr]
|
||||
}
|
||||
|
||||
func (am addrMap) set(addr string, v interface{}) {
|
||||
am[addr] = v
|
||||
}
|
||||
|
||||
func (am addrMap) delete(addr string) {
|
||||
delete(am, addr)
|
||||
}
|
||||
|
||||
type Client struct {
|
||||
cc *ClientConfig
|
||||
c *ssh.Client
|
||||
|
||||
stopc chan sets.Empty
|
||||
stopcEx *sync.Mutex
|
||||
stopcc bool
|
||||
|
||||
wakec chan sets.Empty
|
||||
lfc chan LocalForwardReq
|
||||
rfc chan RemoteForwardReq
|
||||
|
||||
lfclosec chan LocalForwardReq
|
||||
rfclosec chan RemoteForwardReq
|
||||
|
||||
localForwards portMap
|
||||
remoteForwards portMap
|
||||
}
|
||||
|
||||
func NewClient(cc *ClientConfig) *Client {
|
||||
c := &Client{
|
||||
cc: cc,
|
||||
|
||||
stopc: make(chan sets.Empty),
|
||||
stopcEx: &sync.Mutex{},
|
||||
|
||||
wakec: make(chan sets.Empty),
|
||||
lfc: make(chan LocalForwardReq),
|
||||
rfc: make(chan RemoteForwardReq),
|
||||
|
||||
lfclosec: make(chan LocalForwardReq),
|
||||
rfclosec: make(chan RemoteForwardReq),
|
||||
|
||||
localForwards: portMap{},
|
||||
remoteForwards: portMap{},
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *Client) Stop(ctx context.Context) {
|
||||
c.stopcEx.Lock()
|
||||
defer c.stopcEx.Unlock()
|
||||
if !c.stopcc {
|
||||
close(c.stopc)
|
||||
c.stopcc = true
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) Start(ctx context.Context) {
|
||||
pingT := time.NewTimer(17 * time.Second)
|
||||
pingFailCount := 0
|
||||
const pingMaxFail = 3
|
||||
|
||||
const (
|
||||
stateInit = iota
|
||||
stateOK
|
||||
)
|
||||
state := stateInit
|
||||
stateC := make(chan int)
|
||||
stateInitRetryInterval := 7 * time.Second
|
||||
stateInitRetryT := time.NewTicker(stateInitRetryInterval)
|
||||
for {
|
||||
switch state {
|
||||
case stateOK:
|
||||
// check forwards and start
|
||||
case stateInit:
|
||||
if sshc, err := c.connect(ctx); err != nil {
|
||||
log.Errorf("ssh connect: %v", err)
|
||||
} else {
|
||||
c.c = sshc
|
||||
state = stateOK
|
||||
go func() {
|
||||
defer c.c.Conn.Close()
|
||||
|
||||
err := c.c.Conn.Wait()
|
||||
if err != nil {
|
||||
log.Errorf("ssh client conn: %v", err)
|
||||
}
|
||||
select {
|
||||
case stateC <- stateInit:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case req := <-c.lfc:
|
||||
if c.c != nil {
|
||||
c.localForward(ctx, req)
|
||||
}
|
||||
case req := <-c.rfc:
|
||||
if c.c != nil {
|
||||
c.remoteForward(ctx, req)
|
||||
}
|
||||
case req := <-c.lfclosec:
|
||||
c.localForwardClose(ctx, req)
|
||||
case req := <-c.rfclosec:
|
||||
c.remoteForwardClose(ctx, req)
|
||||
case <-c.wakec:
|
||||
break
|
||||
case <-pingT.C:
|
||||
//TODO ping check
|
||||
//ping fail
|
||||
if pingFailCount > pingMaxFail {
|
||||
state = stateInit
|
||||
}
|
||||
case newState := <-stateC:
|
||||
state = newState
|
||||
case <-stateInitRetryT.C:
|
||||
case <-c.stopc:
|
||||
if c.c != nil {
|
||||
c.c.Conn.Close()
|
||||
}
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) connect(ctx context.Context) (*ssh.Client, error) {
|
||||
sshc, err := c.cc.NewClient(ctx)
|
||||
return sshc, err
|
||||
}
|
||||
|
||||
func (c *Client) LocalForward(ctx context.Context, req LocalForwardReq) {
|
||||
select {
|
||||
case c.lfc <- req:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) localForward(ctx context.Context, req LocalForwardReq) {
|
||||
if err := c.localForward_(ctx, req); err != nil {
|
||||
log.Errorf("local forward: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) localForward_(ctx context.Context, req LocalForwardReq) error {
|
||||
// check LocalAddr/LocalPort existence
|
||||
if c.localForwards.contains(req.LocalPort, req.LocalAddr) {
|
||||
return errors.Errorf("local addr occupied: %s:%d", req.LocalAddr, req.LocalPort)
|
||||
}
|
||||
|
||||
addr := net.JoinHostPort(req.LocalAddr, fmt.Sprintf("%d", req.LocalPort))
|
||||
listener, err := net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "tcp listen %s", addr)
|
||||
}
|
||||
fwd := &forwarder{
|
||||
listener: listener,
|
||||
|
||||
dial: c.c.Dial,
|
||||
dialAddr: req.RemoteAddr,
|
||||
dialPort: req.RemotePort,
|
||||
|
||||
done: c.localForwardDone,
|
||||
doneAddr: req.LocalAddr,
|
||||
donePort: req.LocalPort,
|
||||
|
||||
tick: req.Tick,
|
||||
tickCb: req.TickCb,
|
||||
}
|
||||
|
||||
c.localForwards.set(req.LocalPort, req.LocalAddr, fwd)
|
||||
go fwd.Start(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) localForwardDone(laddr string, lport int) {
|
||||
c.localForwards.delete(lport, laddr)
|
||||
}
|
||||
|
||||
func (c *Client) RemoteForward(ctx context.Context, req RemoteForwardReq) {
|
||||
select {
|
||||
case c.rfc <- req:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) remoteForward(ctx context.Context, req RemoteForwardReq) {
|
||||
if err := c.remoteForward_(ctx, req); err != nil {
|
||||
log.Errorf("remote forward: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) remoteForward_(ctx context.Context, req RemoteForwardReq) error {
|
||||
// check RemoteAddr/RemotePort existence
|
||||
if c.remoteForwards.contains(req.RemotePort, req.RemoteAddr) {
|
||||
return errors.Errorf("remote addr occupied: %s:%d", req.RemoteAddr, req.RemotePort)
|
||||
}
|
||||
|
||||
addr := net.JoinHostPort(req.RemoteAddr, fmt.Sprintf("%d", req.RemotePort))
|
||||
listener, err := c.c.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
return errors.Wrapf(err, "ssh listen %s", addr)
|
||||
}
|
||||
|
||||
fwd := &forwarder{
|
||||
listener: listener,
|
||||
|
||||
dial: net.Dial,
|
||||
dialAddr: req.LocalAddr,
|
||||
dialPort: req.LocalPort,
|
||||
|
||||
done: c.remoteForwardDone,
|
||||
doneAddr: req.RemoteAddr,
|
||||
donePort: req.RemotePort,
|
||||
|
||||
tick: req.Tick,
|
||||
tickCb: req.TickCb,
|
||||
}
|
||||
c.remoteForwards.set(req.RemotePort, req.RemoteAddr, fwd)
|
||||
go fwd.Start(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) remoteForwardDone(raddr string, rport int) {
|
||||
c.remoteForwards.delete(rport, raddr)
|
||||
}
|
||||
|
||||
func (c *Client) LocalForwardClose(ctx context.Context, req LocalForwardReq) {
|
||||
select {
|
||||
case c.lfclosec <- req:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) localForwardClose(ctx context.Context, req LocalForwardReq) {
|
||||
v := c.localForwards.get(req.LocalPort, req.LocalAddr)
|
||||
if v != nil {
|
||||
fwd := v.(*forwarder)
|
||||
fwd.Stop(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) RemoteForwardClose(ctx context.Context, req RemoteForwardReq) {
|
||||
select {
|
||||
case c.rfclosec <- req:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) remoteForwardClose(ctx context.Context, req RemoteForwardReq) {
|
||||
v := c.remoteForwards.get(req.RemotePort, req.RemoteAddr)
|
||||
if v != nil {
|
||||
fwd := v.(*forwarder)
|
||||
fwd.Stop(ctx)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
// Copyright 2019 Yunion
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package ssh
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
|
||||
"golang.org/x/crypto/ssh"
|
||||
|
||||
"yunion.io/x/pkg/errors"
|
||||
)
|
||||
|
||||
type ClientConfig struct {
|
||||
User string
|
||||
Host string
|
||||
Port int
|
||||
Key string
|
||||
}
|
||||
|
||||
func (cc *ClientConfig) NewClient(ctx context.Context) (*ssh.Client, error) {
|
||||
signer, err := ssh.ParsePrivateKey([]byte(cc.Key))
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "parse ssh key")
|
||||
}
|
||||
sshcc := &ssh.ClientConfig{
|
||||
User: cc.User,
|
||||
Auth: []ssh.AuthMethod{
|
||||
ssh.PublicKeys(signer),
|
||||
},
|
||||
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
|
||||
}
|
||||
|
||||
addr := net.JoinHostPort(cc.Host, fmt.Sprintf("%d", cc.Port))
|
||||
d := &net.Dialer{}
|
||||
netconn, err := d.DialContext(ctx, "tcp", addr)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "net dial")
|
||||
}
|
||||
|
||||
sshconn, chans, reqs, err := ssh.NewClientConn(netconn, addr, sshcc)
|
||||
if err != nil {
|
||||
netconn.Close()
|
||||
return nil, errors.Wrap(err, "ssh new client conn")
|
||||
}
|
||||
|
||||
sshc := ssh.NewClient(sshconn, chans, reqs)
|
||||
return sshc, nil
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
// Copyright 2019 Yunion
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package ssh
|
||||
|
||||
import (
|
||||
"context"
|
||||
)
|
||||
|
||||
type epClientSet struct {
|
||||
cc ClientConfig
|
||||
clients []*Client
|
||||
|
||||
mark bool
|
||||
}
|
||||
|
||||
func (epcs *epClientSet) clearMark() {
|
||||
epcs.mark = false
|
||||
}
|
||||
|
||||
func (epcs *epClientSet) setMark() {
|
||||
epcs.mark = true
|
||||
}
|
||||
|
||||
func (epcs *epClientSet) getMark() bool {
|
||||
return epcs.mark
|
||||
}
|
||||
|
||||
func (epcs *epClientSet) stop(ctx context.Context) {
|
||||
for _, client := range epcs.clients {
|
||||
client.Stop(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
type epClients map[string]*epClientSet // key: epKey
|
||||
|
||||
type ClientSet struct {
|
||||
epClients epClients
|
||||
}
|
||||
|
||||
func NewClientSet() *ClientSet {
|
||||
cs := &ClientSet{
|
||||
epClients: epClients{},
|
||||
}
|
||||
return cs
|
||||
}
|
||||
|
||||
func (cs *ClientSet) ClearAllMark() {
|
||||
for _, epcs := range cs.epClients {
|
||||
epcs.clearMark()
|
||||
}
|
||||
}
|
||||
|
||||
func (cs *ClientSet) ResetIfChanged(ctx context.Context, epKey string, cc ClientConfig) bool {
|
||||
epcs, ok := cs.epClients[epKey]
|
||||
if ok {
|
||||
if epcs.cc != cc {
|
||||
epcs.stop(ctx)
|
||||
delete(cs.epClients, epKey)
|
||||
return true
|
||||
}
|
||||
epcs.setMark()
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (cs *ClientSet) AddIfNotExist(ctx context.Context, epKey string, cc ClientConfig) bool {
|
||||
epcs, ok := cs.epClients[epKey]
|
||||
if !ok {
|
||||
epcs := &epClientSet{
|
||||
cc: cc,
|
||||
}
|
||||
epcs.setMark()
|
||||
cs.epClients[epKey] = epcs
|
||||
return true
|
||||
}
|
||||
epcs.setMark()
|
||||
return false
|
||||
}
|
||||
|
||||
func (cs *ClientSet) ResetUnmarked(ctx context.Context) {
|
||||
for epKey, epcs := range cs.epClients {
|
||||
if !epcs.getMark() {
|
||||
epcs.stop(ctx)
|
||||
delete(cs.epClients, epKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (cs *ClientSet) ForwardKeySet() ForwardKeySet {
|
||||
fks := ForwardKeySet{}
|
||||
for epKey, epcs := range cs.epClients {
|
||||
for _, client := range epcs.clients {
|
||||
fks.addByPortMap(epKey, ForwardKeyTypeL, client.localForwards)
|
||||
fks.addByPortMap(epKey, ForwardKeyTypeR, client.remoteForwards)
|
||||
}
|
||||
}
|
||||
return fks
|
||||
}
|
||||
|
||||
func (cs *ClientSet) LocalForward(ctx context.Context, epKey string, req LocalForwardReq) {
|
||||
client, created := cs.getOrCreateClient(epKey, ForwardKeyTypeL)
|
||||
if created {
|
||||
go client.Start(ctx)
|
||||
}
|
||||
client.LocalForward(ctx, req)
|
||||
}
|
||||
|
||||
func (cs *ClientSet) RemoteForward(ctx context.Context, epKey string, req RemoteForwardReq) {
|
||||
client, created := cs.getOrCreateClient(epKey, ForwardKeyTypeR)
|
||||
if created {
|
||||
go client.Start(ctx)
|
||||
}
|
||||
client.RemoteForward(ctx, req)
|
||||
}
|
||||
|
||||
func (cs *ClientSet) CloseForward(ctx context.Context, fk ForwardKey) {
|
||||
client := cs.getClient(fk.EpKey, fk.Type)
|
||||
if client == nil {
|
||||
return
|
||||
}
|
||||
switch fk.Type {
|
||||
case ForwardKeyTypeL:
|
||||
client.LocalForwardClose(ctx, LocalForwardReq{
|
||||
LocalAddr: fk.KeyAddr,
|
||||
LocalPort: fk.KeyPort,
|
||||
})
|
||||
case ForwardKeyTypeR:
|
||||
client.RemoteForwardClose(ctx, RemoteForwardReq{
|
||||
RemoteAddr: fk.KeyAddr,
|
||||
RemotePort: fk.KeyPort,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
func (cs *ClientSet) LocalForwardClose(ctx context.Context, epKey string, req LocalForwardReq) {
|
||||
client := cs.getClient(epKey, ForwardKeyTypeL)
|
||||
client.LocalForwardClose(ctx, req)
|
||||
}
|
||||
|
||||
func (cs *ClientSet) RemoteForwardClose(ctx context.Context, epKey string, req RemoteForwardReq) {
|
||||
client := cs.getClient(epKey, ForwardKeyTypeR)
|
||||
client.RemoteForwardClose(ctx, req)
|
||||
}
|
||||
*/
|
||||
|
||||
func (cs *ClientSet) getOrCreateClient(epKey string, typ string) (*Client, bool) {
|
||||
return cs.getClient_(epKey, typ, true)
|
||||
}
|
||||
|
||||
func (cs *ClientSet) getClient(epKey string, typ string) *Client {
|
||||
client, _ := cs.getClient_(epKey, typ, false)
|
||||
return client
|
||||
}
|
||||
|
||||
func (cs *ClientSet) getClient_(epKey string, typ string, create bool) (*Client, bool) {
|
||||
var client *Client
|
||||
|
||||
clients, ok := cs.epClients[epKey]
|
||||
if !ok || len(clients.clients) == 0 {
|
||||
if !create {
|
||||
return nil, false
|
||||
}
|
||||
client = NewClient(&clients.cc)
|
||||
clients.clients = append(clients.clients, client)
|
||||
cs.epClients[epKey] = clients
|
||||
return client, true
|
||||
} else {
|
||||
client = clients.clients[0]
|
||||
return client, false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
// Copyright 2019 Yunion
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package ssh // import "yunion.io/x/onecloud/pkg/proxyagent/agent/ssh"
|
||||
@@ -0,0 +1,69 @@
|
||||
// Copyright 2019 Yunion
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package ssh
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
)
|
||||
|
||||
const (
|
||||
ForwardKeyTypeL = "L"
|
||||
ForwardKeyTypeR = "R"
|
||||
)
|
||||
|
||||
type ForwardKey struct {
|
||||
EpKey string
|
||||
Type string
|
||||
KeyAddr string
|
||||
KeyPort int
|
||||
|
||||
Value interface{}
|
||||
}
|
||||
|
||||
func (fk *ForwardKey) Key() string {
|
||||
return fmt.Sprintf("%s/%s/%s/%d", fk.EpKey, fk.Type, fk.KeyAddr, fk.KeyPort)
|
||||
}
|
||||
|
||||
type ForwardKeySet map[string]ForwardKey
|
||||
|
||||
func (fks ForwardKeySet) addByPortMap(epKey, typ string, pm portMap) {
|
||||
for port, addrMap := range pm {
|
||||
for addr := range addrMap {
|
||||
fk := ForwardKey{
|
||||
EpKey: epKey,
|
||||
Type: typ,
|
||||
KeyAddr: addr,
|
||||
KeyPort: port,
|
||||
}
|
||||
fks.Add(fk)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (fks ForwardKeySet) Contains(fk ForwardKey) bool {
|
||||
key := fk.Key()
|
||||
if _, ok := fks[key]; ok {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (fks ForwardKeySet) Remove(fk ForwardKey) {
|
||||
delete(fks, fk.Key())
|
||||
}
|
||||
|
||||
func (fks ForwardKeySet) Add(fk ForwardKey) {
|
||||
fks[fk.Key()] = fk
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
// Copyright 2019 Yunion
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package ssh
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"yunion.io/x/log"
|
||||
)
|
||||
|
||||
type TickFunc func(context.Context)
|
||||
|
||||
type LocalForwardReq struct {
|
||||
LocalAddr string
|
||||
LocalPort int
|
||||
RemoteAddr string
|
||||
RemotePort int
|
||||
|
||||
Tick time.Duration
|
||||
TickCb TickFunc
|
||||
}
|
||||
|
||||
type RemoteForwardReq struct {
|
||||
// LocalAddr is the address the forward will forward to
|
||||
LocalAddr string
|
||||
// LocalPort is the port the forward will forward to
|
||||
LocalPort int
|
||||
|
||||
// RemoteAddr is the address on the remote to listen on
|
||||
RemoteAddr string
|
||||
// RemotePort is the address on the remote to listen on
|
||||
RemotePort int
|
||||
|
||||
Tick time.Duration
|
||||
TickCb TickFunc
|
||||
}
|
||||
|
||||
type dialFunc func(n, addr string) (net.Conn, error)
|
||||
type doneFunc func(laddr string, lport int)
|
||||
|
||||
type forwarder struct {
|
||||
listener net.Listener
|
||||
|
||||
dial dialFunc
|
||||
dialAddr string
|
||||
dialPort int
|
||||
|
||||
done doneFunc
|
||||
doneAddr string
|
||||
donePort int
|
||||
|
||||
tick time.Duration
|
||||
tickCb TickFunc
|
||||
}
|
||||
|
||||
func (fwd *forwarder) Stop(ctx context.Context) {
|
||||
fwd.listener.Close()
|
||||
}
|
||||
|
||||
func (fwd *forwarder) Start(
|
||||
ctx context.Context,
|
||||
) {
|
||||
var (
|
||||
listener = fwd.listener
|
||||
dial = fwd.dial
|
||||
dialAddr = fwd.dialAddr
|
||||
dialPort = fwd.dialPort
|
||||
done = fwd.done
|
||||
doneAddr = fwd.doneAddr
|
||||
donePort = fwd.donePort
|
||||
tick = fwd.tick
|
||||
tickCb = fwd.tickCb
|
||||
)
|
||||
|
||||
ctx, cancelFunc := context.WithCancel(ctx)
|
||||
|
||||
if done != nil {
|
||||
defer done(doneAddr, donePort)
|
||||
}
|
||||
|
||||
defer listener.Close()
|
||||
|
||||
go func() { // accept local/remote connection
|
||||
for {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
log.Warningf("local forward: accept: %v", err)
|
||||
cancelFunc()
|
||||
break
|
||||
}
|
||||
go func(local net.Conn) {
|
||||
defer local.Close()
|
||||
|
||||
// dial remote/local
|
||||
addr := net.JoinHostPort(dialAddr, fmt.Sprintf("%d", dialPort))
|
||||
remote, err := dial("tcp", addr)
|
||||
if err != nil {
|
||||
log.Warningf("local forward: dial remote: %v", err)
|
||||
return
|
||||
}
|
||||
defer remote.Close()
|
||||
|
||||
// forward
|
||||
go io.Copy(local, remote)
|
||||
go io.Copy(remote, local)
|
||||
<-ctx.Done()
|
||||
}(conn)
|
||||
}
|
||||
}()
|
||||
|
||||
if tick > 0 && tickCb != nil {
|
||||
go func() {
|
||||
ticker := time.NewTicker(tick)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
tickCb(ctx)
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user