Skip to content

Commit cad2c36

Browse files
author
David Capwell
committed
integrated ssh native
1 parent a08e6ef commit cad2c36

File tree

3 files changed

+151
-59
lines changed

3 files changed

+151
-59
lines changed

‎ssh.go

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,12 @@ func createTasks(hosts []string, cmd string, options Options) chan workpool.Task
114114
}
115115

116116
func newSshTask(host string, cmd string, opt Options) workpool.Task {
117+
//TODO find better way to do this. Switching should be a user option, and default based off num hosts
118+
117119
// use this method to switch impls
118-
return newSshProcessTask(host, cmd, opt)
120+
// return newSshProcessTask(host, cmd, opt)
121+
// when num hosts is < 8~, proc seems faster, else native seems faster
122+
return newSshNativeTask(host, cmd, opt)
119123
}
120124

121125
func runConcurrency(options Options, numHosts int) int {
@@ -128,3 +132,12 @@ func runConcurrency(options Options, numHosts int) int {
128132
}
129133
return conc
130134
}
135+
136+
func createContext(host string) SshResponseContext {
137+
rsp := SshResponse{}
138+
ctx := SshResponseContext{
139+
Hostname: host,
140+
Response: rsp,
141+
}
142+
return ctx
143+
}

‎ssh_native.go

Lines changed: 133 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,70 +1,160 @@
11
package gossh
22

33
import (
4-
"fmt"
5-
"os"
6-
"io"
7-
"io/ioutil"
8-
"path/filepath"
9-
"code.google.com/p/go.crypto/ssh"
4+
"bytes"
5+
"code.google.com/p/go.crypto/ssh"
6+
"fmt"
7+
"io"
8+
"io/ioutil"
9+
"os"
10+
"path/filepath"
1011
)
1112

12-
// create ssh task
13+
// create a new native ssh task
14+
func newSshNativeTask(host string, cmd string, opt Options) func() (interface{}, error) {
15+
state := &sshNativeTask{
16+
Host: host,
17+
Cmd: cmd,
18+
Opts: opt,
19+
}
20+
return state.run
21+
}
22+
23+
type sshNativeTask struct {
24+
Host string
25+
Cmd string
26+
Opts Options
27+
}
28+
29+
// workpool task function. Runs ssh cmd
30+
func (s *sshNativeTask) run() (interface{}, error) {
31+
client, err := s.newClient()
32+
// create session
33+
session, err := client.NewSession()
34+
if err != nil {
35+
return nil, err
36+
}
37+
defer session.Close()
38+
39+
// run cmd
40+
var stdout bytes.Buffer
41+
var stderr bytes.Buffer
42+
session.Stdout = &stdout
43+
session.Stderr = &stderr
44+
45+
ctx := createContext(s.Host)
46+
if err := session.Run(s.Cmd); err != nil {
47+
// if of type ExitError, then its a remote issue, add to code
48+
waitmsg, ok := err.(*ssh.ExitError)
49+
if ok {
50+
ctx.Response.Code = waitmsg.ExitStatus()
51+
} else {
52+
// else return err
53+
return ctx, err
54+
}
55+
}
56+
57+
ctx.Response.Stdout = stdout.String()
58+
ctx.Response.Stderr = stderr.String()
59+
60+
return ctx, nil
61+
}
62+
63+
func (s *sshNativeTask) newClient() (*ssh.ClientConn, error) {
64+
config, err := s.newClientConfig()
65+
if err != nil {
66+
return nil, err
67+
}
68+
// create client
69+
//TODO allow user to override port
70+
client, err := ssh.Dial("tcp", s.Host+":22", config)
71+
if err != nil {
72+
return nil, err
73+
}
74+
return client, nil
75+
}
76+
77+
func (s *sshNativeTask) newClientConfig() (*ssh.ClientConfig, error) {
78+
// get auth
79+
// create keychain; only keys are supported right now
80+
var kc *keychain
81+
var err error
82+
if s.Opts.Identity != "" {
83+
kc, err = newKeychainWithKeys(s.Opts.Identity)
84+
} else {
85+
kc, err = newKeychain()
86+
}
87+
if err != nil {
88+
return nil, err
89+
}
90+
// get user
91+
user := s.Opts.User
92+
if user == "" {
93+
user = os.Getenv("USER")
94+
}
95+
config := &ssh.ClientConfig{
96+
User: user,
97+
Auth: []ssh.ClientAuth{
98+
ssh.ClientAuthKeyring(kc),
99+
},
100+
}
101+
return config, nil
102+
}
13103

14104
// pulled from go.crytpo/ssh/test/test_unix_test.go
15105
type keychain struct {
16-
keys []ssh.Signer
106+
keys []ssh.Signer
17107
}
18108

19109
func newKeychain() (kc *keychain, err error) {
20-
kc = new(keychain)
21-
err = kc.load()
22-
return
110+
kc = new(keychain)
111+
err = kc.load()
112+
return
23113
}
24114

25115
func newKeychainWithKeys(keys ...string) (*keychain, error) {
26-
kc := new(keychain)
27-
for _, key := range keys {
28-
err := kc.loadPEM(key)
29-
if err != nil {
30-
return nil, err
31-
}
32-
}
33-
return kc, nil
116+
kc := new(keychain)
117+
for _, key := range keys {
118+
err := kc.loadPEM(key)
119+
if err != nil {
120+
return nil, err
121+
}
122+
}
123+
return kc, nil
34124
}
35125

36126
func (k *keychain) Key(i int) (ssh.PublicKey, error) {
37-
if i < 0 || i >= len(k.keys) {
38-
return nil, nil
39-
}
40-
return k.keys[i].PublicKey(), nil
127+
if i < 0 || i >= len(k.keys) {
128+
return nil, nil
129+
}
130+
return k.keys[i].PublicKey(), nil
41131
}
42132

43133
func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
44-
return k.keys[i].Sign(rand, data)
134+
return k.keys[i].Sign(rand, data)
45135
}
46136

47137
func (k *keychain) load() error {
48-
sshDir := fmt.Sprintf("%s/.ssh", os.Getenv("HOME"))
49-
return filepath.Walk(sshDir, func(path string, info os.FileInfo, err error) error {
50-
if err == nil {
51-
// no error reading file info, so lets see if its a pem
52-
// if not, just skip
53-
k.loadPEM(path)
54-
}
55-
return nil
56-
})
138+
sshDir := fmt.Sprintf("%s/.ssh", os.Getenv("HOME"))
139+
return filepath.Walk(sshDir, func(path string, info os.FileInfo, err error) error {
140+
if err == nil {
141+
// no error reading file info, so lets see if its a pem
142+
// if not, just skip
143+
k.loadPEM(path)
144+
}
145+
return nil
146+
})
57147
}
58148

59149
func (k *keychain) loadPEM(file string) error {
60-
buf, err := ioutil.ReadFile(file)
61-
if err != nil {
62-
return err
63-
}
64-
key, err := ssh.ParsePrivateKey(buf)
65-
if err != nil {
66-
return err
67-
}
68-
k.keys = append(k.keys, key)
69-
return nil
150+
buf, err := ioutil.ReadFile(file)
151+
if err != nil {
152+
return err
153+
}
154+
key, err := ssh.ParsePrivateKey(buf)
155+
if err != nil {
156+
return err
157+
}
158+
k.keys = append(k.keys, key)
159+
return nil
70160
}

‎ssh_proc.go

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,13 @@ import (
88
"time"
99
)
1010

11-
func newSshProcessTask(host string, cmd string, opt Options) func()(interface{}, error) {
12-
state := &sshProcessTask{
11+
func newSshProcessTask(host string, cmd string, opt Options) func() (interface{}, error) {
12+
state := &sshProcessTask{
1313
Host: host,
1414
Cmd: cmd,
1515
Options: opt,
1616
}
17-
return state.Run
17+
return state.run
1818
}
1919

2020
// this file is for SshTask that calls the local ssh shell command.
@@ -24,7 +24,7 @@ type sshProcessTask struct {
2424
Options Options
2525
}
2626

27-
func (s *sshProcessTask) Run() (interface{}, error) {
27+
func (s *sshProcessTask) run() (interface{}, error) {
2828
start := time.Now()
2929
// must return of type (SshResponseContext, error)
3030
//cmd := exec.Command("/usr/bin/ssh", s.Host, s.Cmd)
@@ -93,16 +93,6 @@ func (s *sshProcessTask) generateCmdArguments() []string {
9393
return cmd
9494
}
9595

96-
97-
func createContext(host string) SshResponseContext {
98-
rsp := SshResponse{}
99-
ctx := SshResponseContext{
100-
Hostname: host,
101-
Response: rsp,
102-
}
103-
return ctx
104-
}
105-
10696
func exitCode(err error) (int, error) {
10797
if err != nil {
10898
// it puts exit code in err... grrr
@@ -120,4 +110,3 @@ func exitCode(err error) (int, error) {
120110
// was successful, so exit code is 0
121111
return 0, nil
122112
}
123-

0 commit comments

Comments
 (0)