From 0a2068d7fbd447de1135d3b16e9f26de5dd18d56 Mon Sep 17 00:00:00 2001
From: Unknwon <joe2010xtmf@163.com>
Date: Sun, 1 Feb 2015 17:21:56 -0500
Subject: [PATCH] models: code fix on #818

---
 cmd/fix.go                   |  2 +-
 conf/locale/locale_en-US.ini |  4 +--
 models/publickey.go          | 52 +++++++++++++++++++++++-------------
 3 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/cmd/fix.go b/cmd/fix.go
index 5122ff32c7..eff85d6282 100644
--- a/cmd/fix.go
+++ b/cmd/fix.go
@@ -164,7 +164,7 @@ func runFixLocation(ctx *cli.Context) {
 	fmt.Scanln()
 
 	// Fix in authorized_keys file.
-	sshPath := path.Join(models.SshPath, "authorized_keys")
+	sshPath := path.Join(models.SSHPath, "authorized_keys")
 	if com.IsFile(sshPath) {
 		fmt.Printf("Fixing pathes in file: %s\n", sshPath)
 		if err := rewriteAuthorizedKeys(sshPath, oldPath, execPath); err != nil {
diff --git a/conf/locale/locale_en-US.ini b/conf/locale/locale_en-US.ini
index c0979d1bdb..660fb253bb 100644
--- a/conf/locale/locale_en-US.ini
+++ b/conf/locale/locale_en-US.ini
@@ -514,8 +514,8 @@ dashboard.delete_repo_archives = Delete all repositories archives
 dashboard.delete_repo_archives_success = All repositories archives have been deleted successfully.
 dashboard.git_gc_repos = Do garbage collection on repositories
 dashboard.git_gc_repos_success = All repositories have done garbage collection successfully.
-dashboard.resync_all_sshkeys = Do resync .ssh/autorized_key file
-dashboard.resync_all_sshkeys_success = All keys are synced again.
+dashboard.resync_all_sshkeys = Rewrite '.ssh/autorized_key' file(caution: non-Gogs keys will be lost)
+dashboard.resync_all_sshkeys_success = All public keys have been rewritten successfully.
 dashboard.server_uptime = Server Uptime
 dashboard.current_goroutine = Current Goroutines
 dashboard.current_memory_usage = Current Memory Usage
diff --git a/models/publickey.go b/models/publickey.go
index 41233d0c39..67ab4242f2 100644
--- a/models/publickey.go
+++ b/models/publickey.go
@@ -33,7 +33,7 @@ const (
 )
 
 var (
-	ErrKeyAlreadyExist = errors.New("Public key already exist")
+	ErrKeyAlreadyExist = errors.New("Public key already exists")
 	ErrKeyNotExist     = errors.New("Public key does not exist")
 	ErrKeyUnableVerify = errors.New("Unable to verify public key")
 )
@@ -41,7 +41,7 @@ var (
 var sshOpLocker = sync.Mutex{}
 
 var (
-	SshPath string // SSH directory.
+	SSHPath string // SSH directory.
 	appPath string // Execution(binary) path.
 )
 
@@ -72,9 +72,9 @@ func init() {
 	appPath = strings.Replace(appPath, "\\", "/", -1)
 
 	// Determine and create .ssh path.
-	SshPath = filepath.Join(homeDir(), ".ssh")
-	if err = os.MkdirAll(SshPath, 0700); err != nil {
-		log.Fatal(4, "fail to create SshPath(%s): %v\n", SshPath, err)
+	SSHPath = filepath.Join(homeDir(), ".ssh")
+	if err = os.MkdirAll(SSHPath, 0700); err != nil {
+		log.Fatal(4, "fail to create '%s': %v", SSHPath, err)
 	}
 }
 
@@ -248,12 +248,13 @@ func saveAuthorizedKeyFile(keys ...*PublicKey) error {
 	sshOpLocker.Lock()
 	defer sshOpLocker.Unlock()
 
-	fpath := filepath.Join(SshPath, "authorized_keys")
+	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()
+
 	finfo, err := f.Stat()
 	if err != nil {
 		return err
@@ -270,8 +271,7 @@ func saveAuthorizedKeyFile(keys ...*PublicKey) error {
 	}
 
 	for _, key := range keys {
-		_, err = f.WriteString(key.GetAuthorizedString())
-		if err != nil {
+		if _, err = f.WriteString(key.GetAuthorizedString()); err != nil {
 			return err
 		}
 	}
@@ -418,8 +418,8 @@ func DeletePublicKey(key *PublicKey) error {
 		return err
 	}
 
-	fpath := filepath.Join(SshPath, "authorized_keys")
-	tmpPath := filepath.Join(SshPath, "authorized_keys.tmp")
+	fpath := filepath.Join(SSHPath, "authorized_keys")
+	tmpPath := filepath.Join(SSHPath, "authorized_keys.tmp")
 	if err = rewriteAuthorizedKeys(key, fpath, tmpPath); err != nil {
 		return err
 	} else if err = os.Remove(fpath); err != nil {
@@ -428,20 +428,36 @@ func DeletePublicKey(key *PublicKey) error {
 	return os.Rename(tmpPath, fpath)
 }
 
-// RewriteAllPublicKeys remove any authorized key and re-write all key from database again
+// RewriteAllPublicKeys removes any authorized key and rewrite all keys from database again.
 func RewriteAllPublicKeys() error {
-	keys := make([]*PublicKey, 0, 5)
-	err := x.Find(&keys)
+	sshOpLocker.Lock()
+	defer sshOpLocker.Unlock()
+
+	tmpPath := filepath.Join(SSHPath, "authorized_keys.tmp")
+	f, err := os.Create(tmpPath)
+	if err != nil {
+		return err
+	}
+	defer os.Remove(tmpPath)
+
+	err = x.Iterate(new(PublicKey), func(idx int, bean interface{}) (err error) {
+		_, err = f.WriteString((bean.(*PublicKey)).GetAuthorizedString())
+		return err
+	})
+	f.Close()
 	if err != nil {
 		return err
 	}
 
-	fpath := filepath.Join(SshPath, "authorized_keys")
-	if _, err := os.Stat(fpath); os.IsNotExist(err) {
-		return saveAuthorizedKeyFile(keys...)
+	fpath := filepath.Join(SSHPath, "authorized_keys")
+	if com.IsExist(fpath) {
+		if err = os.Remove(fpath); err != nil {
+			return err
+		}
 	}
-	if err := os.Remove(fpath); err != nil {
+	if err = os.Rename(tmpPath, fpath); err != nil {
 		return err
 	}
-	return saveAuthorizedKeyFile(keys...)
+
+	return nil
 }