gowiki/userdb.go

// Copyright (C) 2017 Marius Schellenberger

package main

import (
	"bytes"
	"encoding/gob"
	"errors"
	"github.com/boltdb/bolt"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/crypto/sha3"
)

type User struct {
	Username string `json:"username"`
	Password string `json:"password"`
	Admin    bool   `json"admin"`
	Locked   int    `json:"locked"`
}

func validatePassword(pw string) (b []byte) {
	b = []byte(pw)
	if len(pw) <= 56 {
		return
	}
	hash := sha3.New384()
	hash.Write(b)
	b = hash.Sum(nil)
	return
}

const (
	dbFile       = "gowiki.db"
	defUser      = "admin"
	defPassword  = "gowiki"
	userTable    = "user"
	snippetTable = "snippet"
	bcryptCost   = 13
)

var (
	errUserNotExists = dbErr("user does not exist")
	errUserExists    = dbErr("user already exist")
)

func dbErr(i interface{}) error {
	return fmt.Errorf("db: %s", i)
}

type BoltStore struct {
	Marshaler
	db *bolt.DB
}

func NewBoltStore() (bs *BoltStore, err error) {
	db, err := bolt.Open(dbFile, 0666, nil)
	if err != nil {
		return
	}
	if err = db.Update(func(tx *bolt.Tx) error {
		if _, err := tx.CreateBucketIfNotExists([]byte(userTable)); err != nil {
			return err
		}
		if _, err := tx.CreateBucketIfNotExists([]byte(snippetTable)); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return
	}
	bs = &BoltStore{NewGOB(), db}
	if !bs.UserExists(defUser) {
		err = bs.CreateUser(defUser, defPassword, true)
		if err != nil {
			return
		}
	}
	bs.UnlockUser(defUser)
	return
}

func (bs *BoltStore) Close() error {
	return bs.db.Close()
}

func (bs *BoltStore) GetUsers() (users []User, err error) {
	users, err = bs.DumpUsers()
	for i := range users {
		users[i].Password = ""
	}
	return
}

func (bs *BoltStore) GetUserWithoutPassword(username string) (u User, err error) {
	u, err = bs.GetUser(username)
	u.Password = ""
	return
}

func (bs *BoltStore) GetUser(username string) (u User, err error) {
	err = bs.db.View(func(tx *bolt.Tx) error {
		v := tx.Bucket([]byte(userTable)).Get([]byte(username))
		if v == nil {
			return errUserNotExists
		}
		return bs.Unmarshal(v, &u)
	})
	return
}

func (bs *BoltStore) UserIsAdmin(username string) (admin bool) {
	bs.db.View(func(tx *bolt.Tx) error {
		u, _ := bs.GetUser(username)
		admin = u.Admin
		return nil
	})
	return
}

func (bs *BoltStore) UserExists(username string) error {
	return bs.db.View(func(tx *bolt.Tx) error {
		if tx.Bucket([]byte(userTable)).Get([]byte(username)) == nil {
			return errUserNotExists
		}
		return nil
	})
}

func (bs *BoltStore) DumpUsers() (users []User, err error) {
	var buf [][]byte
	err = bs.db.View(func(tx *bolt.Tx) error {
		return tx.Bucket([]byte(userTable)).ForEach(func(k, v []byte) error {
			buf = append(buf, copyBuf(v))
			return nil
		})
	})
	if err != nil {
		return
	}
	for _, b := range buf {
		u := new(User)
		if err = bs.Unmarshal(b, u); err != nil {
			return
		}
		users = append(users, *u)
	}
	return
}

func (bs *BoltStore) RestoreUsers(users []User) error {
	return bs.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(userTable))
		err := b.ForEach(func(k, v []byte) error {
			return b.Delete(k)
		})
		if err != nil {
			return err
		}
		for _, u := range users {
			v, err := bs.Marshal(u)
			if err != nil {
				return err
			}
			err = b.Put([]byte(u.Username), v)
			if err != nil {
				return err
			}
		}
		return nil
	})
}

func (bs *BoltStore) CreateUser(username, password string, admin bool) error {
	if err := bs.UserExists(username); err == nil {
		return errUserExists
	}
	return bs.db.Update(func(tx *bolt.Tx) error {
		hash, err := bcrypt.GenerateFromPassword(validatePassword(password), bcryptCost)
		v, err := bs.Marshal(&User{
			Username: username,
			Password: string(hash),
			Admin:    admin,
		})
		if err != nil {
			return err
		}
		return tx.Bucket([]byte(userTable)).Put([]byte(username), v)
	})
}

func (bs *BoltStore) Login(username, password string) error {
	var e error
	err := bs.db.Update(func(tx *bolt.Tx) error {
		us, err := bs.GetUser(username)
		if err != nil {
			return err
		}
		if us.Locked == 3 {
			return errors.New("db: user locked")
		}
		b := tx.Bucket([]byte(userTable))
		if err := bcrypt.CompareHashAndPassword([]byte(us.Password), validatePassword(password)); err != nil {
			us.Locked++
			v, err := bs.Marshal(us)
			if err != nil {
				return err
			}
			if err = b.Put([]byte(username), v); err != nil {
				return err
			}
			e = errors.New("db: wrong login")
			return nil
		}
		us.Locked = 0
		v, err := bs.Marshal(us)
		if err != nil {
			return err
		}
		return b.Put([]byte(username), v)
	})
	if e != nil {
		return e
	}
	return err
}

func (bs *BoltStore) UnlockUser(username string) error {
	return bs.db.Update(func(tx *bolt.Tx) error {
		u, err := bs.GetUser(username)
		if err != nil {
			return err
		}
		if u.Locked < 3 {
			return nil
		}
		u.Locked = 0
		v, err := bs.Marshal(u)
		if err != nil {
			return err
		}
		return tx.Bucket([]byte(userTable)).Put([]byte(username), v)
	})
}

func (bs *BoltStore) UpdateUserPassword(username, password string) error {
	return bs.db.Update(func(tx *bolt.Tx) error {
		if password != "" {
			us, err := bs.GetUser(username)
			if err != nil {
				return err
			}
			hash, err := bcrypt.GenerateFromPassword(validatePassword(password), bcryptCost)
			if err != nil {
				return err
			}
			us.Password = string(hash)
			v, err := bs.Marshal(us)
			if err != nil {
				return err
			}
			return tx.Bucket([]byte(userTable)).Put([]byte(username), v)
		}
		return nil
	})
}

func (bs *BoltStore) AdminUpdateUser(username, password string, admin bool) error {
	if username == defUser && !admin {
		return errors.New("user '" + defUser + "' cannot lose admin privileges")
	}
	return bs.db.Update(func(tx *bolt.Tx) error {
		us, err := bs.GetUser(username)
		if err != nil {
			return err
		}
		us.Admin = admin
		if password != "" {
			hash, err := bcrypt.GenerateFromPassword(validatePassword(password), bcryptCost)
			if err != nil {
				return err
			}
			us.Password = string(hash)
		}
		v, err := bs.Marshal(us)
		if err != nil {
			return err
		}
		return tx.Bucket([]byte(userTable)).Put([]byte(username), v)
	})
}

func (bs *BoltStore) DeleteUser(username string) error {
	if username == defUser {
		return errors.New("user '" + defUser + "' cannot be deleted")
	}
	if err := bs.UserExists(username); err != nil {
		return err
	}
	return bs.db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket([]byte(userTable)).Delete([]byte(username))
	})
}

type Marshaler interface {
	Marshal(v interface{}) ([]byte, error)
	Unmarshal(data []byte, v interface{}) error
}

type gobMarshaler struct{}

func NewGOB() Marshaler {
	return gobMarshaler{}
}

func (gobMarshaler) Marshal(v interface{}) (b []byte, err error) {
	buf := new(bytes.Buffer)
	err = gob.NewEncoder(buf).Encode(v)
	b = buf.Bytes()
	return
}

func (gobMarshaler) Unmarshal(data []byte, v interface{}) error {
	return gob.NewDecoder(bytes.NewBuffer(data)).Decode(v)
}