view auth/persistent.go @ 200:8426a92fda00

Persistent session store: Implmented Renew.
author Sascha L. Teichmann <teichmann@intevation.de>
date Sun, 22 Jul 2018 09:25:25 +0200
parents ddc7ef95c247
children 80dc7bbe97db
line wrap: on
line source

package auth

import (
	"bytes"
	"database/sql"
	"log"
	"time"

	bolt "github.com/coreos/bbolt"
)

type PersistentConnectionPool struct {
	db    *bolt.DB
	conns map[string]*Connection
	cmds  chan func(*PersistentConnectionPool)
}

var sessionsBucket = []byte("sessions")

func NewPersistentConnectionPool(filename string) (*PersistentConnectionPool, error) {

	db, err := bolt.Open(filename, 0600, nil)
	if err != nil {
		return nil, err
	}

	conns := make(map[string]*Connection)
	err = db.Update(func(tx *bolt.Tx) error {
		b, err := tx.CreateBucketIfNotExists(sessionsBucket)
		if err != nil {
			return err
		}

		// pre-load sessions
		c := b.Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			var conn Connection
			if err := conn.deserialize(bytes.NewReader(v)); err != nil {
				return err
			}
			conns[string(k)] = &conn
		}

		return nil
	})

	if err != nil {
		db.Close()
		return nil, err
	}

	pcp := &PersistentConnectionPool{
		db:    db,
		conns: conns,
		cmds:  make(chan func(*PersistentConnectionPool)),
	}
	go pcp.run()
	return pcp, nil
}

func (pcp *PersistentConnectionPool) run() {
	for {
		select {
		case cmd := <-pcp.cmds:
			cmd(pcp)
		case <-time.After(time.Minute):
			pcp.cleanDB()
		case <-time.After(time.Minute * 5):
			pcp.cleanToken()
		}
	}
}

func (pcp *PersistentConnectionPool) cleanDB() {
	log.Println("cleanDB: Not implemented, yet.")
}

func (pcp *PersistentConnectionPool) cleanToken() {
	log.Println("cleanToken: Not implemented, yet.")
}

func (pcp *PersistentConnectionPool) remove(token string) {
	err := pcp.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(sessionsBucket)
		return b.Delete([]byte(token))
	})
	if err != nil {
		log.Printf("error: %v\n", err)
	}
}

func (pcp *PersistentConnectionPool) Delete(token string) bool {
	res := make(chan bool)
	pcp.cmds <- func(pcp *PersistentConnectionPool) {
		conn, found := pcp.conns[token]
		if !found {
			res <- false
			return
		}
		conn.close()
		delete(pcp.conns, token)
		pcp.remove(token)
		res <- true
	}
	return <-res
}

func (pcp *PersistentConnectionPool) store(token string, con *Connection) {
	err := pcp.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(sessionsBucket)
		var buf bytes.Buffer
		if err := con.serialize(&buf); err != nil {
			return err
		}
		return b.Put([]byte(token), buf.Bytes())
	})
	if err != nil {
		log.Printf("error: %v\n", err)
	}
}

func (pcp *PersistentConnectionPool) Add(token string, session *Session) *Connection {
	res := make(chan *Connection)

	pcp.cmds <- func(cp *PersistentConnectionPool) {
		con := pcp.conns[token]
		if con == nil {
			con = &Connection{}
			pcp.conns[token] = con
		}
		con.set(session)
		pcp.store(token, con)
		res <- con
	}

	con := <-res
	return con
}

func (pcp *PersistentConnectionPool) Renew(token string) (string, error) {

	type result struct {
		newToken string
		err      error
	}

	resCh := make(chan result)

	pcp.cmds <- func(cp *PersistentConnectionPool) {
		con := pcp.conns[token]
		if con == nil {
			resCh <- result{err: ErrNoSuchToken}
		} else {
			delete(pcp.conns, token)
			pcp.remove(token)
			newToken := GenerateSessionKey()
			// TODO: Ensure that this is not racy!
			con.session.ExpiresAt = time.Now().Add(maxTokenValid).Unix()
			pcp.conns[newToken] = con
			pcp.store(newToken, con)
			resCh <- result{newToken: newToken}
		}
	}

	r := <-resCh
	return r.newToken, r.err
}

func (pcp *PersistentConnectionPool) Do(token string, fn func(*sql.DB) error) error {
	log.Println("Do: Not implemented, yet.")
	return nil
}

func (pcp *PersistentConnectionPool) Session(token string) *Session {
	log.Println("Session: Not implemented, yet.")
	return nil
}

func (pcp *PersistentConnectionPool) Shutdown() error {
	log.Println("info: shutdown persistent connection pool.")
	if db := pcp.db; db != nil {
		pcp.db = nil
		return db.Close()
	}
	return nil
}