alibaba/kt-connect

package transmission

import (
    "fmt"
    "github.com/alibaba/kt-connect/pkg/common"
    "github.com/alibaba/kt-connect/pkg/kt/service/sshchannel"
    "github.com/alibaba/kt-connect/pkg/kt/util"
    "github.com/rs/zerolog/log"
    "strings"
    "time"
)

// ForwardPodToLocal mapping pod port to local port
func ForwardPodToLocal(exposePorts, podName, privateKey string) (int, error) {
    log.Info().Msgf("Forwarding pod %s to local via port %s", podName, exposePorts)
    localSshPort := util.GetRandomTcpPort()

    // port forward pod 22 -> local <random port>
    if _, err := SetupPortForwardToLocal(podName, common.StandardSshPort, localSshPort); err != nil {
        return -1, err
    }

    err := ForwardRemotePortsViaSshTunnel(exposePorts, localSshPort, privateKey)
    if err != nil {
        return -1, err
    }

    return localSshPort, nil
}

// ForwardRemotePortsViaSshTunnel forward multiple remote ports to local
func ForwardRemotePortsViaSshTunnel(exposePorts string, localSshPort int, privateKey string) error {
    // supports multi port-pairs
    portPairs := strings.Split(exposePorts, ",")
    res := make(chan error)
    for _, exposePort := range portPairs {
        localPort, remotePort, err2 := util.ParsePortMapping(exposePort)
        if err2 != nil {
            return err2
        }
        forwardRemotePortViaSshTunnel(localPort, remotePort, localSshPort, privateKey, res)
    }
    select {
    case err := <-res:
        return err
    case <-time.After(1 * time.Second):
        go func() {
            // consume the res channel to avoid block reverse tunnel
            <-res
        }()
    }
    return nil
}

// ForwardRemotePortViaSshTunnel forward remote pod to local
func forwardRemotePortViaSshTunnel(localPort, remotePort, localSshPort int, privateKey string, res chan error) {
    remoteEndpoint := fmt.Sprintf("127.0.0.1:%d", localSshPort)
    localEndpoint := fmt.Sprintf("0.0.0.0:%d", remotePort)
    sshAddress := fmt.Sprintf("127.0.0.1:%d", localPort)
    log.Debug().Msgf("Forwarding %s to local endpoint %s via %s", remoteEndpoint, localEndpoint, sshAddress)
    sshReverseTunnel(privateKey, remoteEndpoint, localEndpoint, sshAddress, res)
}

func sshReverseTunnel(privateKey, remoteEndpoint, localEndpoint, sshAddress string, res chan error) {
    go func() {
        err := sshchannel.Ins().ForwardRemoteToLocal(privateKey, remoteEndpoint, localEndpoint, sshAddress)
        if err != nil {
            if res != nil {
                log.Error().Err(err).Msgf("Failed to setup reverse tunnel")
                res <-err
            } else {
                log.Debug().Err(err).Msgf("Reverse tunnel interrupted")
            }
        }

        time.Sleep(10 * time.Second)
        log.Debug().Msgf("Reverse tunnel reconnecting ...")
        sshReverseTunnel(privateKey, remoteEndpoint, localEndpoint, sshAddress, nil)
    }()
}