diff --git a/conn_pool.go b/conn_pool.go index fb530b9..6e244b0 100644 --- a/conn_pool.go +++ b/conn_pool.go @@ -20,30 +20,35 @@ func generateToken() (string, error) { var ErrSessionExpired = errors.New("session expired") +type Session struct { + imapConn *imapclient.Client + username, password string +} + // TODO: expiration timer type ConnPool struct { locker sync.Mutex - conns map[string]*imapclient.Client + sessions map[string]*Session } func NewConnPool() *ConnPool { return &ConnPool{ - conns: make(map[string]*imapclient.Client), + sessions: make(map[string]*Session), } } -func (pool *ConnPool) Get(token string) (*imapclient.Client, error) { +func (pool *ConnPool) Get(token string) (*Session, error) { pool.locker.Lock() defer pool.locker.Unlock() - conn, ok := pool.conns[token] + session, ok := pool.sessions[token] if !ok { return nil, ErrSessionExpired } - return conn, nil + return session, nil } -func (pool *ConnPool) Put(conn *imapclient.Client) (token string, err error) { +func (pool *ConnPool) Put(imapConn *imapclient.Client, username, password string) (token string, err error) { pool.locker.Lock() defer pool.locker.Unlock() @@ -51,22 +56,26 @@ func (pool *ConnPool) Put(conn *imapclient.Client) (token string, err error) { var err error token, err = generateToken() if err != nil { - conn.Logout() + imapConn.Logout() return "", err } - if _, ok := pool.conns[token]; !ok { + if _, ok := pool.sessions[token]; !ok { break } } - pool.conns[token] = conn + pool.sessions[token] = &Session{ + imapConn: imapConn, + username: username, + password: password, + } go func() { - <-conn.LoggedOut() + <-imapConn.LoggedOut() pool.locker.Lock() - delete(pool.conns, token) + delete(pool.sessions, token) pool.locker.Unlock() }() diff --git a/server.go b/server.go index 6e800d0..be174b8 100644 --- a/server.go +++ b/server.go @@ -93,6 +93,7 @@ func NewServer(imapURL, smtpURL string) (*Server, error) { type context struct { echo.Context server *Server + session *Session conn *imapclient.Client } @@ -126,7 +127,7 @@ func handleLogin(ectx echo.Context) error { return ctx.Render(http.StatusOK, "login.html", nil) } - token, err := ctx.server.imap.pool.Put(conn) + token, err := ctx.server.imap.pool.Put(conn, username, password) if err != nil { return fmt.Errorf("failed to put connection in pool: %v", err) } @@ -167,6 +168,8 @@ func handleGetPart(ctx *context, raw bool) error { disp, dispParams, _ := part.Header.ContentDisposition() filename := dispParams["filename"] + // TODO: set Content-Length if possible + if !strings.EqualFold(mimeType, "text/plain") || strings.EqualFold(disp, "attachment") { dispParams := make(map[string]string) if filename != "" { @@ -235,13 +238,14 @@ func New(imapURL, smtpURL string) *echo.Echo { return err } - ctx.conn, err = ctx.server.imap.pool.Get(cookie.Value) + ctx.session, err = ctx.server.imap.pool.Get(cookie.Value) if err == ErrSessionExpired { ctx.setToken("") return ctx.Redirect(http.StatusFound, "/login") } else if err != nil { return err } + ctx.conn = ctx.session.imapConn return next(ctx) }