Skip to content

Commit

Permalink
feat(lang): Support git credentials (#145)
Browse files Browse the repository at this point in the history
* WIP

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* feat(CLI): Support ssh agent forwarding in up

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Update

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
  • Loading branch information
gaocegege committed May 16, 2022
1 parent 9d9d282 commit 0804dc7
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 12 deletions.
1 change: 1 addition & 0 deletions pkg/lang/ir/graph.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ func NewGraph() *Graph {
// They are used by vscode remote.
"curl",
"openssh-client",
"git",
},

PyPIPackages: []string{},
Expand Down
50 changes: 38 additions & 12 deletions pkg/ssh/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import (
"github.com/sirupsen/logrus"
"github.com/tensorchord/envd/pkg/lang/ir"
"golang.org/x/crypto/ssh"
"golang.org/x/crypto/ssh/agent"
"golang.org/x/term"
)

Expand All @@ -40,8 +41,7 @@ type Client interface {
}

type generalClient struct {
config *ssh.ClientConfig
server string
cli *ssh.Client
}

func NewClient(server, user string,
Expand All @@ -54,6 +54,8 @@ func NewClient(server, user string,
},
}

var cli *ssh.Client

if auth {
// read private key file
pemBytes, err := ioutil.ReadFile(privateKeyPath)
Expand All @@ -70,27 +72,51 @@ func NewClient(server, user string,
}
}

host := fmt.Sprintf("%s:%d", server, port)
// open connection
conn, err := ssh.Dial("tcp", host, config)
if err != nil {
return nil, errors.Wrap(err, "dialing failed")
}
cli = conn

// open connection to the local agent
socketLocation := os.Getenv("SSH_AUTH_SOCK")
if socketLocation != "" {
agentConn, err := net.Dial("unix", socketLocation)
if err != nil {
return nil, errors.Wrap(err, "could not connect to local agent socket")
}
// create agent and add in auth
forwardingAgent := agent.NewClient(agentConn)
// add callback for forwarding agent to SSH config
// XXX - might want to handle reconnects appending multiple callbacks
auth := ssh.PublicKeysCallback(forwardingAgent.Signers)
config.Auth = append(config.Auth, auth)
if err := agent.ForwardToAgent(cli, forwardingAgent); err != nil {
return nil, errors.Wrap(err, "forwarding agent to client failed")
}
}

return &generalClient{
config: config,
server: fmt.Sprintf("%v:%v", server, port),
cli: cli,
}, nil
}

func (c generalClient) Attach() error {
// open connection
conn, err := ssh.Dial("tcp", c.server, c.config)
if err != nil {
return fmt.Errorf("dial to %v failed %v", c.server, err)
}
defer conn.Close()
defer c.cli.Close()

// open session
session, err := conn.NewSession()
session, err := c.cli.NewSession()
if err != nil {
return fmt.Errorf("create session for %v failed %v", c.server, err)
return errors.Wrap(err, "creating session failed")
}
defer session.Close()

if err := agent.RequestAgentForwarding(session); err != nil {
return errors.Wrap(err, "requesting agent forwarding failed")
}

modes := ssh.TerminalModes{
ssh.ECHO: 0, // Disable echoing
ssh.ECHOCTL: 0, // Don't print control chars
Expand Down

0 comments on commit 0804dc7

Please sign in to comment.