Change SSHTunnelList to struct to make Open() semantics better.

(cherry picked from commit 48f672af92)
This commit is contained in:
CJ Cullen 2015-06-16 10:36:38 -07:00 committed by Brendan Burns
parent e98f79e4bc
commit 4cd4d363c5
2 changed files with 32 additions and 23 deletions

View File

@ -210,7 +210,7 @@ type Master struct {
InsecureHandler http.Handler
// Used for secure proxy
tunnels util.SSHTunnelList
tunnels *util.SSHTunnelList
tunnelsLock sync.Mutex
installSSHKey InstallSSHKey
}
@ -772,7 +772,7 @@ func (m *Master) Dial(net, addr string) (net.Conn, error) {
}
func (m *Master) needToReplaceTunnels(addrs []string) bool {
if len(m.tunnels) != len(addrs) {
if m.tunnels == nil || m.tunnels.Len() != len(addrs) {
return true
}
// TODO (cjcullen): This doesn't need to be n^2
@ -850,7 +850,7 @@ func (m *Master) setupSecureProxy(user, keyfile string) {
if err := m.loadTunnels(user, keyfile); err != nil {
glog.Errorf("Failed to load SSH Tunnels: %v", err)
}
if len(m.tunnels) != 0 {
if m.tunnels != nil && m.tunnels.Len() != 0 {
// Sleep for 10 seconds if we have some tunnels.
// TODO (cjcullen): tunnels can lag behind actually existing nodes.
time.Sleep(9 * time.Second)

View File

@ -207,9 +207,11 @@ type SSHTunnelEntry struct {
Tunnel *SSHTunnel
}
type SSHTunnelList []SSHTunnelEntry
type SSHTunnelList struct {
entries []SSHTunnelEntry
}
func MakeSSHTunnels(user, keyfile string, addresses []string) (SSHTunnelList, error) {
func MakeSSHTunnels(user, keyfile string, addresses []string) (*SSHTunnelList, error) {
tunnels := []SSHTunnelEntry{}
for ix := range addresses {
addr := addresses[ix]
@ -219,19 +221,22 @@ func MakeSSHTunnels(user, keyfile string, addresses []string) (SSHTunnelList, er
}
tunnels = append(tunnels, SSHTunnelEntry{addr, tunnel})
}
return tunnels, nil
return &SSHTunnelList{tunnels}, nil
}
func (l SSHTunnelList) Open() error {
for ix := 0; ix < len(l); ix++ {
if err := l[ix].Tunnel.Open(); err != nil {
// Remove a failed Open from the list.
glog.Errorf("Failed to open tunnel %v: %v", l[ix], err)
l = append(l[:ix], l[ix+1:]...)
ix--
// Open attempts to open all tunnels in the list, and removes any tunnels that
// failed to open.
func (l *SSHTunnelList) Open() error {
var openTunnels []SSHTunnelEntry
for ix := range l.entries {
if err := l.entries[ix].Tunnel.Open(); err != nil {
glog.Errorf("Failed to open tunnel %v: %v", l.entries[ix], err)
} else {
openTunnels = append(openTunnels, l.entries[ix])
}
}
if len(l) == 0 {
l.entries = openTunnels
if len(l.entries) == 0 {
return errors.New("Failed to open any tunnels.")
}
return nil
@ -240,9 +245,9 @@ func (l SSHTunnelList) Open() error {
// Close asynchronously closes all tunnels in the list after waiting for 1
// minute. Tunnels will still be open upon this function's return, but should
// no longer be used.
func (l SSHTunnelList) Close() {
for ix := range l {
entry := l[ix]
func (l *SSHTunnelList) Close() {
for ix := range l.entries {
entry := l.entries[ix]
go func() {
defer HandleCrash()
time.Sleep(1 * time.Minute)
@ -253,22 +258,26 @@ func (l SSHTunnelList) Close() {
}
}
func (l SSHTunnelList) Dial(network, addr string) (net.Conn, error) {
if len(l) == 0 {
func (l *SSHTunnelList) Dial(network, addr string) (net.Conn, error) {
if len(l.entries) == 0 {
return nil, fmt.Errorf("Empty tunnel list.")
}
return l[mathrand.Int()%len(l)].Tunnel.Dial(network, addr)
return l.entries[mathrand.Int()%len(l.entries)].Tunnel.Dial(network, addr)
}
func (l SSHTunnelList) Has(addr string) bool {
for ix := range l {
if l[ix].Address == addr {
func (l *SSHTunnelList) Has(addr string) bool {
for ix := range l.entries {
if l.entries[ix].Address == addr {
return true
}
}
return false
}
func (l *SSHTunnelList) Len() int {
return len(l.entries)
}
func EncodePrivateKey(private *rsa.PrivateKey) []byte {
return pem.EncodeToMemory(&pem.Block{
Bytes: x509.MarshalPKCS1PrivateKey(private),