LevyraatiBot/auth.go

package main

import (
	"container/heap"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"log"
	"net/http"
	"time"
)

type SessionData struct {
	sid      string
	username string
	priority time.Time
	index    int
}

type SessionQueue []*SessionData

func (pq SessionQueue) Len() int {
	return len(pq)
}

func (pq SessionQueue) Less(i, j int) bool {
	return pq[i].priority.Before(pq[j].priority)
}

func (pq SessionQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

func (pq *SessionQueue) Push(x interface{}) {
	n := len(*pq)
	item := x.(*SessionData)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *SessionQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

var (
	sessions     = make(map[string]*SessionData)
	sessionQueue = make(SessionQueue, 0)
)

func sessionExpirer() {
	for {
		for len(sessionQueue) > 0 && time.Now().After(sessionQueue[0].priority) {
			session := heap.Pop(&sessionQueue).(*SessionData)
			delete(sessions, session.sid)
		}
		time.Sleep(time.Second * 5)
	}
}

func init() {
	go sessionExpirer()
}

func addSession(data *SessionData) {
	sessions[data.sid] = data
	heap.Push(&sessionQueue, data)
}

func userOk(username, password string) bool {
	return (username == "Lamperi" && password == "paskaa")
}

func tryLogin(username, password string, longerTime bool) (http.Cookie, error) {
	if exists := userOk(username, password); !exists {
		return http.Cookie{},
			errors.New("The username or password you entered isn't correct.")
	}

	sid, err := randString(32)
	if err != nil {
		return http.Cookie{}, err
	}

	duration := time.Hour * 1
	if longerTime {
		duration = time.Hour * 24 * 14
	}

	loginCookie := http.Cookie{
		Name:     "id",
		Value:    sid,
		MaxAge:   int(duration.Seconds()),
		HttpOnly: true,
	}

	expiration := time.Now().Add(duration)
	addSession(&SessionData{sid, username, expiration, 0})

	return loginCookie, nil
}

func getSession(req *http.Request) (*SessionData, error) {
	cookie, err := req.Cookie("id")
	if err != nil {
		return nil, err
	}

	session, exists := sessions[cookie.Value]
	if !exists {
		return nil, errors.New("Session expired from server")
	}
	return session, nil
}

func randString(size int) (string, error) {
	buf := make([]byte, size)

	if _, err := rand.Read(buf); err != nil {
		log.Println(err)
		return "", errors.New("Couldn't generate random string")
	}

	return base64.URLEncoding.EncodeToString(buf)[:size], nil
}