Improve delete SSH key

4 files changed, 39 insertions, 59 deletions

diff --git a/models/issue.go b/models/issue.go
index 6a32312eed..64fc45a61d 100644
--- a/models/issue.go
+++ b/models/issue.go
@@ -30,19 +30,17 @@ type Issue struct {
 	IsPull          bool // Indicates whether is a pull request or not.
 	IsClosed        bool
 	Labels          string `xorm:"TEXT"`
-	Mentions        string `xorm:"TEXT"`
 	Content         string `xorm:"TEXT"`
 	RenderedContent string `xorm:"-"`
+	Priority        int
 	NumComments     int
+	Deadline        time.Time
 	Created         time.Time `xorm:"created"`
 	Updated         time.Time `xorm:"updated"`
 }
 
 // CreateIssue creates new issue for repository.
 func CreateIssue(userId, repoId, milestoneId, assigneeId int64, issueCount int, name, labels, content string, isPull bool) (issue *Issue, err error) {
-	// TODO: find out mentions
-	mentions := ""
-
 	sess := orm.NewSession()
 	defer sess.Close()
 	sess.Begin()
@@ -56,7 +54,6 @@ func CreateIssue(userId, repoId, milestoneId, assigneeId int64, issueCount int,
 		AssigneeId:  assigneeId,
 		IsPull:      isPull,
 		Labels:      labels,
-		Mentions:    mentions,
 		Content:     content,
 	}
 	if _, err = sess.Insert(issue); err != nil {
diff --git a/models/models.go b/models/models.go
index ebb558fcc8..6e4f7d1022 100644
--- a/models/models.go
+++ b/models/models.go
@@ -139,10 +139,9 @@ func NewEngine() (err error) {
 
 type Statistic struct {
 	Counter struct {
-		User, PublicKey, Repo,
-		Watch, Action, Access,
-		Issue, Comment,
-		Mirror, Oauth, Release int64
+		User, PublicKey, Repo, Watch, Action, Access,
+		Issue, Comment, Mirror, Oauth, Release,
+		LoginSource, Webhook int64
 	}
 }
 
@@ -158,6 +157,8 @@ func GetStatistic() (stats Statistic) {
 	stats.Counter.Mirror, _ = orm.Count(new(Mirror))
 	stats.Counter.Oauth, _ = orm.Count(new(Oauth2))
 	stats.Counter.Release, _ = orm.Count(new(Release))
+	stats.Counter.LoginSource, _ = orm.Count(new(LoginSource))
+	stats.Counter.Webhook, _ = orm.Count(new(Webhook))
 	return
 }
 
diff --git a/models/publickey.go b/models/publickey.go
index b80412812b..e594bbe91b 100644
--- a/models/publickey.go
+++ b/models/publickey.go
@@ -6,12 +6,11 @@ package models
 
 import (
 	"bufio"
+	"bytes"
 	"errors"
 	"fmt"
-	"io"
 	"io/ioutil"
 	"os"
-	"os/exec"
 	"path"
 	"path/filepath"
 	"strings"
@@ -19,7 +18,9 @@ import (
 	"time"
 
 	"github.com/Unknwon/com"
+	qlog "github.com/qiniu/log"
 
+	"github.com/gogits/gogs/modules/base"
 	"github.com/gogits/gogs/modules/log"
 )
 
@@ -30,29 +31,21 @@ const (
 
 var (
 	ErrKeyAlreadyExist = errors.New("Public key already exist")
+	ErrKeyNotExist     = errors.New("Public key does not exist")
 )
 
 var sshOpLocker = sync.Mutex{}
 
 var (
-	sshPath string
-	appPath string
+	sshPath string // SSH directory.
+	appPath string // Execution(binary) path.
 )
 
-// exePath returns the executable path.
-func exePath() (string, error) {
-	file, err := exec.LookPath(os.Args[0])
-	if err != nil {
-		return "", err
-	}
-	return filepath.Abs(file)
-}
-
 // homeDir returns the home directory of current user.
 func homeDir() string {
 	home, err := com.HomeDir()
 	if err != nil {
-		return "/"
+		qlog.Fatalln(err)
 	}
 	return home
 }
@@ -60,17 +53,14 @@ func homeDir() string {
 func init() {
 	var err error
 
-	appPath, err = exePath()
-	if err != nil {
-		fmt.Printf("publickey.init(fail to get app path): %v\n", err)
-		os.Exit(2)
+	if appPath, err = base.ExecDir(); err != nil {
+		qlog.Fatalf("publickey.init(fail to get app path): %v\n", err)
 	}
 
 	// Determine and create .ssh path.
 	sshPath = filepath.Join(homeDir(), ".ssh")
 	if err = os.MkdirAll(sshPath, os.ModePerm); err != nil {
-		fmt.Printf("publickey.init(fail to create sshPath(%s)): %v\n", sshPath, err)
-		os.Exit(2)
+		qlog.Fatalf("publickey.init(fail to create sshPath(%s)): %v\n", sshPath, err)
 	}
 }
 
@@ -129,8 +119,8 @@ func AddPublicKey(key *PublicKey) (err error) {
 	return nil
 }
 
+// rewriteAuthorizedKeys finds and deletes corresponding line in authorized_keys file.
 func rewriteAuthorizedKeys(key *PublicKey, p, tmpP string) error {
-	// Delete SSH key in SSH key file.
 	sshOpLocker.Lock()
 	defer sshOpLocker.Unlock()
 
@@ -146,55 +136,48 @@ func rewriteAuthorizedKeys(key *PublicKey, p, tmpP string) error {
 	}
 	defer fw.Close()
 
-	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
-			}
+	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
 		}
 
 		// Found the line and copy rest of file.
-		if strings.Contains(line, fmt.Sprintf("key-%d", key.Id)) && strings.Contains(line, key.Content) {
+		if !isFound && bytes.Contains(line, keyword) && bytes.Contains(line, content) {
+			isFound = true
 			continue
 		}
+
 		// Still finding the line, copy the line that currently read.
-		if _, err = fw.WriteString(line + "\n"); err != nil {
+		if _, err = fw.Write(line); err != nil {
 			return err
 		}
-
-		if errRead == io.EOF {
-			break
-		}
 	}
+
 	return nil
 }
 
 // DeletePublicKey deletes SSH key information both in database and authorized_keys file.
-func DeletePublicKey(key *PublicKey) (err error) {
-	// Delete SSH key in database.
-	has, err := orm.Id(key.Id).Get(key)
+func DeletePublicKey(key *PublicKey) error {
+	has, err := orm.Get(key)
 	if err != nil {
 		return err
 	} else if !has {
-		return errors.New("Public key does not exist")
+		return ErrKeyNotExist
 	}
+
 	if _, err = orm.Delete(key); err != nil {
 		return err
 	}
 
 	p := filepath.Join(sshPath, "authorized_keys")
 	tmpP := filepath.Join(sshPath, "authorized_keys.tmp")
-	log.Trace("ssh.DeletePublicKey(authorized_keys): %s", p)
+	log.Trace("publickey.DeletePublicKey(authorized_keys): %s", p)
 
 	if err = rewriteAuthorizedKeys(key, p, tmpP); err != nil {
 		return err
diff --git a/models/user.go b/models/user.go
index a4b753f8e1..720065c3f1 100644
--- a/models/user.go
+++ b/models/user.go
@@ -32,7 +32,6 @@ var (
 	ErrUserNotExist     = errors.New("User does not exist")
 	ErrEmailAlreadyUsed = errors.New("E-mail already used")
 	ErrUserNameIllegal  = errors.New("User name contains illegal characters")
-	ErrKeyNotExist      = errors.New("Public key does not exist")
 )
 
 // User represents the object of individual and member of organization.
@@ -315,12 +314,12 @@ func DeleteUser(user *User) error {
 	}
 
 	// Delete all SSH keys.
-	keys := make([]PublicKey, 0, 10)
+	keys := make([]*PublicKey, 0, 10)
 	if err = orm.Find(&keys, &PublicKey{OwnerId: user.Id}); err != nil {
 		return err
 	}
 	for _, key := range keys {
-		if err = DeletePublicKey(&key); err != nil {
+		if err = DeletePublicKey(key); err != nil {
 			return err
 		}
 	}