1 files changed, 61 insertions, 52 deletions

diff --git a/models/publickey.go b/models/publickey.go
index e594bbe91b..7ef25b9ceb 100644
--- a/models/publickey.go
+++ b/models/publickey.go
@@ -6,9 +6,9 @@ package models
 
 import (
 	"bufio"
-	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"path"
@@ -26,7 +26,7 @@ import (
 
 const (
 	// "### autogenerated by gitgos, DO NOT EDIT\n"
-	TPL_PUBLICK_KEY = `command="%s serv key-%d",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty %s`
+	_TPL_PUBLICK_KEY = `command="%s serv key-%d",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty %s` + "\n"
 )
 
 var (
@@ -64,7 +64,7 @@ func init() {
 	}
 }
 
-// PublicKey represents a SSH key of user.
+// PublicKey represents a SSH key.
 type PublicKey struct {
 	Id          int64
 	OwnerId     int64  `xorm:"UNIQUE(s) INDEX NOT NULL"`
@@ -75,14 +75,29 @@ type PublicKey struct {
 	Updated     time.Time `xorm:"UPDATED"`
 }
 
-// GenAuthorizedKey returns formatted public key string.
-func GenAuthorizedKey(keyId int64, key string) string {
-	return fmt.Sprintf(TPL_PUBLICK_KEY+"\n", appPath, keyId, key)
+// GetAuthorizedString generates and returns formatted public key string for authorized_keys file.
+func (key *PublicKey) GetAuthorizedString() string {
+	return fmt.Sprintf(_TPL_PUBLICK_KEY, appPath, key.Id, key.Content)
 }
 
-// AddPublicKey adds new public key to database and SSH key file.
+// saveAuthorizedKeyFile writes SSH key content to authorized_keys file.
+func saveAuthorizedKeyFile(key *PublicKey) error {
+	sshOpLocker.Lock()
+	defer sshOpLocker.Unlock()
+
+	fpath := filepath.Join(sshPath, "authorized_keys")
+	f, err := os.OpenFile(fpath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	_, err = f.WriteString(key.GetAuthorizedString())
+	return err
+}
+
+// AddPublicKey adds new public key to database and authorized_keys file.
 func AddPublicKey(key *PublicKey) (err error) {
-	// Check if public key name has been used.
 	has, err := orm.Get(key)
 	if err != nil {
 		return err
@@ -91,7 +106,7 @@ func AddPublicKey(key *PublicKey) (err error) {
 	}
 
 	// Calculate fingerprint.
-	tmpPath := strings.Replace(filepath.Join(os.TempDir(), fmt.Sprintf("%d", time.Now().Nanosecond()),
+	tmpPath := strings.Replace(path.Join(os.TempDir(), fmt.Sprintf("%d", time.Now().Nanosecond()),
 		"id_rsa.pub"), "\\", "/", -1)
 	os.MkdirAll(path.Dir(tmpPath), os.ModePerm)
 	if err = ioutil.WriteFile(tmpPath, []byte(key.Content), os.ModePerm); err != nil {
@@ -108,8 +123,8 @@ func AddPublicKey(key *PublicKey) (err error) {
 	// Save SSH key.
 	if _, err = orm.Insert(key); err != nil {
 		return err
-	}
-	if err = SaveAuthorizedKeyFile(key); err != nil {
+	} else if err = saveAuthorizedKeyFile(key); err != nil {
+		// Roll back.
 		if _, err2 := orm.Delete(key); err2 != nil {
 			return err2
 		}
@@ -119,6 +134,13 @@ func AddPublicKey(key *PublicKey) (err error) {
 	return nil
 }
 
+// ListPublicKey returns a list of all public keys that user has.
+func ListPublicKey(uid int64) ([]PublicKey, error) {
+	keys := make([]PublicKey, 0, 5)
+	err := orm.Find(&keys, &PublicKey{OwnerId: uid})
+	return keys, err
+}
+
 // rewriteAuthorizedKeys finds and deletes corresponding line in authorized_keys file.
 func rewriteAuthorizedKeys(key *PublicKey, p, tmpP string) error {
 	sshOpLocker.Lock()
@@ -137,28 +159,38 @@ func rewriteAuthorizedKeys(key *PublicKey, p, tmpP string) error {
 	defer fw.Close()
 
 	isFound := false
-	keyword := []byte(fmt.Sprintf("key-%d", key.Id))
-	content := []byte(key.Content)
-
-	snr := bufio.NewScanner(fr)
-	for snr.Scan() {
-		line := append(bytes.TrimSpace(snr.Bytes()), '\n')
-		if len(line) == 0 {
-			continue
+	keyword := fmt.Sprintf("key-%d", key.Id)
+	buf := bufio.NewReader(fr)
+	for {
+		line, errRead := buf.ReadString('\n')
+		line = strings.TrimSpace(line)
+
+		if errRead != nil {
+			if errRead != io.EOF {
+				return errRead
+			}
+
+			// Reached end of file, if nothing to read then break,
+			// otherwise handle the last line.
+			if len(line) == 0 {
+				break
+			}
 		}
 
 		// Found the line and copy rest of file.
-		if !isFound && bytes.Contains(line, keyword) && bytes.Contains(line, content) {
+		if !isFound && strings.Contains(line, keyword) && strings.Contains(line, key.Content) {
 			isFound = true
 			continue
 		}
-
 		// Still finding the line, copy the line that currently read.
-		if _, err = fw.Write(line); err != nil {
+		if _, err = fw.WriteString(line + "\n"); err != nil {
 			return err
 		}
-	}
 
+		if errRead == io.EOF {
+			break
+		}
+	}
 	return nil
 }
 
@@ -175,37 +207,14 @@ func DeletePublicKey(key *PublicKey) error {
 		return err
 	}
 
-	p := filepath.Join(sshPath, "authorized_keys")
-	tmpP := filepath.Join(sshPath, "authorized_keys.tmp")
-	log.Trace("publickey.DeletePublicKey(authorized_keys): %s", p)
+	fpath := filepath.Join(sshPath, "authorized_keys")
+	tmpPath := filepath.Join(sshPath, "authorized_keys.tmp")
+	log.Trace("publickey.DeletePublicKey(authorized_keys): %s", fpath)
 
-	if err = rewriteAuthorizedKeys(key, p, tmpP); err != nil {
+	if err = rewriteAuthorizedKeys(key, fpath, tmpPath); err != nil {
 		return err
-	} else if err = os.Remove(p); err != nil {
+	} else if err = os.Remove(fpath); err != nil {
 		return err
 	}
-	return os.Rename(tmpP, p)
-}
-
-// ListPublicKey returns a list of public keys that user has.
-func ListPublicKey(userId int64) ([]PublicKey, error) {
-	keys := make([]PublicKey, 0)
-	err := orm.Find(&keys, &PublicKey{OwnerId: userId})
-	return keys, err
-}
-
-// SaveAuthorizedKeyFile writes SSH key content to SSH key file.
-func SaveAuthorizedKeyFile(key *PublicKey) error {
-	sshOpLocker.Lock()
-	defer sshOpLocker.Unlock()
-
-	p := filepath.Join(sshPath, "authorized_keys")
-	f, err := os.OpenFile(p, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-
-	_, err = f.WriteString(GenAuthorizedKey(key.Id, key.Content))
-	return err
+	return os.Rename(tmpPath, fpath)
 }