4 files changed, 17 insertions, 11 deletions
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)
 }