Change rsa.PublicKey to ssh.PublicKey in Dryad Prepare

Users will use SSH keys rather than plain RSA keys, so Dryad interface
should accept SSH public key instead of generating it from RSA public
key.

Change-Id: I6e757199a7e8a0d3258c1c17ac0eee8412f2b415
Signed-off-by: Maciej Wereski <m.wereski@partner.samsung.com>

10 files changed, 43 insertions, 27 deletions

diff --git a/boruta.go b/boruta.go
index 2f87ea5..cde8065 100644
--- a/boruta.go
+++ b/boruta.go
@@ -29,6 +29,8 @@ import (
 	"net"
 	"strconv"
 	"time"
+
+	"golang.org/x/crypto/ssh"
 )
 
 // ReqState denotes state of the Request.
@@ -229,10 +231,10 @@ type Dryad interface {
 	// PutInMaintenance prepares MuxPi for administrative action.
 	// It blinks LEDs, prints msg on the OLED display, etc.
 	PutInMaintenance(msg string) (err error)
-	// Prepare creates appropriate user, generates RSA key, installs public key
+	// Prepare creates appropriate user, generates SSH keypair, installs public key
 	// so that it can be used for SSH authentication and returns private key.
 	// It removes current instance of the user, etc.
-	Prepare(key *rsa.PublicKey) (err error)
+	Prepare(key *ssh.PublicKey) (err error)
 	// Healthcheck tests Dryad for system state, STM functions and state on MuxPi.
 	// It may cause Dryad to call SetFail of Worker interface if the problem detected is critical.
 	Healthcheck() (err error)
diff --git a/dryad/key.go b/dryad/key.go
index 3881709..9a41b74 100644
--- a/dryad/key.go
+++ b/dryad/key.go
@@ -17,6 +17,7 @@
 package dryad
 
 import (
+	"errors"
 	"os"
 	"path"
 	"strconv"
@@ -25,7 +26,10 @@ import (
 )
 
 // installPublicKey marshals and stores key in a proper location to be read by ssh daemon.
-func installPublicKey(key ssh.PublicKey, homedir, uid, gid string) error {
+func installPublicKey(key *ssh.PublicKey, homedir, uid, gid string) error {
+	if key == nil {
+		return errors.New("empty public key")
+	}
 	sshDir := path.Join(homedir, ".ssh")
 	err := os.MkdirAll(sshDir, 0755)
 	if err != nil {
@@ -41,7 +45,7 @@ func installPublicKey(key ssh.PublicKey, homedir, uid, gid string) error {
 	if err != nil {
 		return err
 	}
-	_, err = f.Write(ssh.MarshalAuthorizedKey(key))
+	_, err = f.Write(ssh.MarshalAuthorizedKey(*key))
 	return err
 }
 
diff --git a/dryad/rusalka.go b/dryad/rusalka.go
index 7549308..e4af51c 100644
--- a/dryad/rusalka.go
+++ b/dryad/rusalka.go
@@ -21,7 +21,6 @@ package dryad
 
 import (
 	"context"
-	"crypto/rsa"
 	"fmt"
 
 	. "git.tizen.org/tools/boruta"
@@ -62,7 +61,7 @@ func (r *Rusalka) PutInMaintenance(msg string) error {
 }
 
 // Prepare is part of implementation of Dryad interface. Call to Prepare stops LED blinking.
-func (r *Rusalka) Prepare(key *rsa.PublicKey) (err error) {
+func (r *Rusalka) Prepare(key *ssh.PublicKey) (err error) {
 	// Stop maintenance.
 	if r.cancelMaintenance != nil {
 		r.cancelMaintenance()
@@ -82,10 +81,7 @@ func (r *Rusalka) Prepare(key *rsa.PublicKey) (err error) {
 	if err != nil {
 		return fmt.Errorf("user information update failed: %s", err)
 	}
-	// Prepare SSH access (it can't fail as key is of type rsa.PublicKey).
-	sshPubKey, _ := ssh.NewPublicKey(key)
-	// TODO: use ssh.PublicKey instead.
-	return r.dryadUser.generateAndInstallKey(sshPubKey)
+	return r.dryadUser.installKey(key)
 }
 
 // Healthcheck is part of implementation of Dryad interface.
diff --git a/dryad/rusalka_test.go b/dryad/rusalka_test.go
index f3799af..69136ed 100644
--- a/dryad/rusalka_test.go
+++ b/dryad/rusalka_test.go
@@ -23,6 +23,7 @@ import (
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
+	"errors"
 	"os"
 	"os/user"
 	"time"
@@ -32,6 +33,7 @@ import (
 	gomock "github.com/golang/mock/gomock"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
+	"golang.org/x/crypto/ssh"
 )
 
 var _ = Describe("Rusalka", func() {
@@ -84,9 +86,13 @@ var _ = Describe("Rusalka", func() {
 			Skip("must be run as root")
 		}
 
+		err = d.Prepare(nil)
+		Expect(err).To(Equal(errors.New("empty public key")))
 		key, err := rsa.GenerateKey(rand.Reader, 1024)
 		Expect(err).ToNot(HaveOccurred())
-		err = d.Prepare(&key.PublicKey)
+		pubKey, err := ssh.NewPublicKey(&key.PublicKey)
+		Expect(err).ToNot(HaveOccurred())
+		err = d.Prepare(&pubKey)
 		Expect(err).ToNot(HaveOccurred())
 		Expect(sshDir).To(BeADirectory())
 		Expect(authorizedKeysFile).To(BeARegularFile())
diff --git a/dryad/user.go b/dryad/user.go
index ef80812..76cef73 100644
--- a/dryad/user.go
+++ b/dryad/user.go
@@ -137,8 +137,8 @@ func (bu *borutaUser) update() (err error) {
 	return
 }
 
-// generateAndInstallKey calls generateAndInstallKey with parameters retrieved from the user field
+// installKey calls installPublicKey with parameters retrieved from the user field
 // of borutaUser structure. This filed must be set before call to this function by update() method.
-func (bu *borutaUser) generateAndInstallKey(key ssh.PublicKey) error {
+func (bu *borutaUser) installKey(key *ssh.PublicKey) error {
 	return installPublicKey(key, bu.user.HomeDir, bu.user.Uid, bu.user.Gid)
 }
diff --git a/mocks/mock_workers.go b/mocks/mock_workers.go
index d64b584..fee6108 100644
--- a/mocks/mock_workers.go
+++ b/mocks/mock_workers.go
@@ -119,15 +119,15 @@ func (m *MockSuperviser) EXPECT() *MockSuperviserMockRecorder {
 }
 
 // Register mocks base method
-func (m *MockSuperviser) Register(arg0 boruta.Capabilities) error {
-	ret := m.ctrl.Call(m, "Register", arg0)
+func (m *MockSuperviser) Register(arg0 boruta.Capabilities, arg1, arg2 string) error {
+	ret := m.ctrl.Call(m, "Register", arg0, arg1, arg2)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // Register indicates an expected call of Register
-func (mr *MockSuperviserMockRecorder) Register(arg0 interface{}) *gomock.Call {
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Register", reflect.TypeOf((*MockSuperviser)(nil).Register), arg0)
+func (mr *MockSuperviserMockRecorder) Register(arg0, arg1, arg2 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Register", reflect.TypeOf((*MockSuperviser)(nil).Register), arg0, arg1, arg2)
 }
 
 // SetFail mocks base method
diff --git a/rpc/dryad/dryad.go b/rpc/dryad/dryad.go
index 2ee7df3..d6fd053 100644
--- a/rpc/dryad/dryad.go
+++ b/rpc/dryad/dryad.go
@@ -3,8 +3,8 @@
 package dryad
 
 import (
-	"crypto/rsa"
 	. "git.tizen.org/tools/boruta"
+	"golang.org/x/crypto/ssh"
 	"net/rpc"
 )
 
@@ -40,7 +40,7 @@ func (s *DryadService) PutInMaintenance(request *DryadPutInMaintenanceRequest, r
 
 // DryadPrepareRequest is a helper structure for Prepare method.
 type DryadPrepareRequest struct {
-	Key *rsa.PublicKey
+	Key *ssh.PublicKey
 }
 
 // DryadPrepareResponse is a helper structure for Prepare method.
@@ -97,7 +97,7 @@ func (_c *DryadClient) PutInMaintenance(msg string) (err error) {
 }
 
 // Prepare is part of implementation of Dryad calling corresponding method on RPC server.
-func (_c *DryadClient) Prepare(key *rsa.PublicKey) (err error) {
+func (_c *DryadClient) Prepare(key *ssh.PublicKey) (err error) {
 	_request := &DryadPrepareRequest{key}
 	_response := &DryadPrepareResponse{}
 	err = _c.client.Call("Dryad.Prepare", _request, _response)
diff --git a/workers/dryadclientmanager_mock_test.go b/workers/dryadclientmanager_mock_test.go
index fb70ae3..577a25a 100644
--- a/workers/dryadclientmanager_mock_test.go
+++ b/workers/dryadclientmanager_mock_test.go
@@ -4,8 +4,8 @@
 package workers
 
 import (
-	rsa "crypto/rsa"
 	gomock "github.com/golang/mock/gomock"
+	ssh "golang.org/x/crypto/ssh"
 	net "net"
 	reflect "reflect"
 )
@@ -70,7 +70,7 @@ func (mr *MockDryadClientManagerMockRecorder) Healthcheck() *gomock.Call {
 }
 
 // Prepare mocks base method
-func (m *MockDryadClientManager) Prepare(arg0 *rsa.PublicKey) error {
+func (m *MockDryadClientManager) Prepare(arg0 *ssh.PublicKey) error {
 	ret := m.ctrl.Call(m, "Prepare", arg0)
 	ret0, _ := ret[0].(error)
 	return ret0
diff --git a/workers/worker_list_test.go b/workers/worker_list_test.go
index 9264206..8c1f368 100644
--- a/workers/worker_list_test.go
+++ b/workers/worker_list_test.go
@@ -392,7 +392,7 @@ var _ = Describe("WorkerList", func() {
 					It("should work to SetState", func() {
 						gomock.InOrder(
 							dcm.EXPECT().Create(info.dryad),
-							dcm.EXPECT().Prepare(gomock.AssignableToTypeOf(&rsa.PublicKey{})).Return(nil),
+							dcm.EXPECT().Prepare(gomock.Any()).Return(nil),
 							dcm.EXPECT().Close(),
 						)
 
@@ -405,7 +405,7 @@ var _ = Describe("WorkerList", func() {
 					It("should fail to SetState if dryadClientManager fails to prepare client", func() {
 						gomock.InOrder(
 							dcm.EXPECT().Create(info.dryad),
-							dcm.EXPECT().Prepare(gomock.AssignableToTypeOf(&rsa.PublicKey{})).Return(testerr),
+							dcm.EXPECT().Prepare(gomock.Any()).Return(testerr),
 							dcm.EXPECT().Close(),
 						)
 
@@ -785,10 +785,13 @@ var _ = Describe("WorkerList", func() {
 			})
 
 			It("should work to set and get information", func() {
+				// There's only 1 setter and 3 getters, so only 1st getter is checked in loop.
 				for i, set := range setters {
 					get := getters[i]
 					get(wl, worker, set(wl, worker, nil), nil)
 				}
+				getDryad(wl, worker, *dryadAddr, nil)
+				getSSH(wl, worker, *sshdAddr, nil)
 			})
 		})
 		Describe("PrepareWorker", func() {
@@ -855,7 +858,7 @@ var _ = Describe("WorkerList", func() {
 				It("should set worker into IDLE state and prepare a key", func() {
 					gomock.InOrder(
 						dcm.EXPECT().Create(info.dryad),
-						dcm.EXPECT().Prepare(gomock.AssignableToTypeOf(&rsa.PublicKey{})).Return(nil),
+						dcm.EXPECT().Prepare(gomock.Any()).Return(nil),
 						dcm.EXPECT().Close(),
 					)
 
@@ -868,7 +871,7 @@ var _ = Describe("WorkerList", func() {
 				It("should fail to prepare worker if dryadClientManager fails to prepare client", func() {
 					gomock.InOrder(
 						dcm.EXPECT().Create(info.dryad),
-						dcm.EXPECT().Prepare(gomock.AssignableToTypeOf(&rsa.PublicKey{})).Return(testerr),
+						dcm.EXPECT().Prepare(gomock.Any()).Return(testerr),
 						dcm.EXPECT().Close(),
 					)
 
diff --git a/workers/workers.go b/workers/workers.go
index b77a01a..c38ea8a 100644
--- a/workers/workers.go
+++ b/workers/workers.go
@@ -27,6 +27,7 @@ import (
 
 	. "git.tizen.org/tools/boruta"
 	"git.tizen.org/tools/boruta/rpc/dryad"
+	"golang.org/x/crypto/ssh"
 )
 
 // UUID denotes a key in Capabilities where WorkerUUID is stored.
@@ -461,7 +462,11 @@ func (wl *WorkerList) prepareKey(worker WorkerUUID) error {
 	if err != nil {
 		return err
 	}
-	err = client.Prepare(&key.PublicKey)
+	pubKey, err := ssh.NewPublicKey(&key.PublicKey)
+	if err != nil {
+		return err
+	}
+	err = client.Prepare(&pubKey)
 	if err != nil {
 		return err
 	}