3 files changed, 85 insertions, 43 deletions

diff --git a/worker/imap/worker.go b/worker/imap/worker.go
index ceff34d..9fbaf0c 100644
--- a/worker/imap/worker.go
+++ b/worker/imap/worker.go
@@ -4,7 +4,6 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
-	"log"
 	"net/url"
 	"strings"
 
@@ -23,9 +22,6 @@ type imapClient struct {
 }
 
 type IMAPWorker struct {
-	messages chan types.WorkerMessage
-	actions  chan types.WorkerMessage
-
 	config struct {
 		scheme   string
 		insecure bool
@@ -33,33 +29,18 @@ type IMAPWorker struct {
 		user     *url.Userinfo
 	}
 
+	worker  *types.Worker
 	client  *imapClient
 	updates chan client.Update
-	logger  *log.Logger
 }
 
-func NewIMAPWorker(logger *log.Logger) *IMAPWorker {
+func NewIMAPWorker(worker *types.Worker) *IMAPWorker {
 	return &IMAPWorker{
-		messages: make(chan types.WorkerMessage, 50),
-		actions:  make(chan types.WorkerMessage, 50),
-		updates:  make(chan client.Update, 50),
-		logger:   logger,
+		worker:  worker,
+		updates: make(chan client.Update, 50),
 	}
 }
 
-func (w *IMAPWorker) GetMessages() chan types.WorkerMessage {
-	return w.messages
-}
-
-func (w *IMAPWorker) PostAction(msg types.WorkerMessage) {
-	w.actions <- msg
-}
-
-func (w *IMAPWorker) postMessage(msg types.WorkerMessage) {
-	w.logger.Printf("=> %T\n", msg)
-	w.messages <- msg
-}
-
 func (w *IMAPWorker) verifyPeerCert(msg types.WorkerMessage) func(
 	rawCerts [][]byte, _ [][]*x509.Certificate) error {
 
@@ -77,9 +58,9 @@ func (w *IMAPWorker) verifyPeerCert(msg types.WorkerMessage) func(
 			Message:  types.RespondTo(msg),
 			CertPool: pool,
 		}
-		w.postMessage(request)
+		w.worker.PostMessage(request, nil)
 
-		response := <-w.actions
+		response := <-w.worker.Actions
 		if response.InResponseTo() != request {
 			return fmt.Errorf("Expected UI to answer cert request")
 		}
@@ -176,24 +157,24 @@ func (w *IMAPWorker) handleMessage(msg types.WorkerMessage) error {
 func (w *IMAPWorker) Run() {
 	for {
 		select {
-		case msg := <-w.actions:
-			w.logger.Printf("<= %T\n", msg)
+		case msg := <-w.worker.Actions:
+			msg = w.worker.ProcessAction(msg)
 			if err := w.handleMessage(msg); err == errUnsupported {
-				w.postMessage(types.Unsupported{
+				w.worker.PostMessage(types.Unsupported{
 					Message: types.RespondTo(msg),
-				})
+				}, nil)
 			} else if err != nil {
-				w.postMessage(types.Error{
+				w.worker.PostMessage(types.Error{
 					Message: types.RespondTo(msg),
 					Error:   err,
-				})
+				}, nil)
 			} else {
-				w.postMessage(types.Ack{
+				w.worker.PostMessage(types.Ack{
 					Message: types.RespondTo(msg),
-				})
+				}, nil)
 			}
 		case update := <-w.updates:
-			w.logger.Printf("[= %T", update)
+			w.worker.Logger.Printf("(= %T", update)
 		}
 	}
 }
diff --git a/worker/types/worker.go b/worker/types/worker.go
new file mode 100644
index 0000000..a99d432
--- /dev/null
+++ b/worker/types/worker.go
@@ -0,0 +1,59 @@
+package types
+
+import (
+	"log"
+)
+
+type Backend interface {
+	Run()
+}
+
+type Worker struct {
+	Actions   chan WorkerMessage
+	Backend   Backend
+	Callbacks map[WorkerMessage]func(msg WorkerMessage)
+	Messages  chan WorkerMessage
+	Logger    *log.Logger
+}
+
+func (worker *Worker) PostAction(msg WorkerMessage,
+	cb func(msg WorkerMessage)) {
+
+	worker.Logger.Printf("=> %T\n", msg)
+	worker.Actions <- msg
+
+	if cb != nil {
+		worker.Callbacks[msg] = cb
+	}
+}
+
+func (worker *Worker) PostMessage(msg WorkerMessage,
+	cb func(msg WorkerMessage)) {
+
+	worker.Logger.Printf("-> %T\n", msg)
+	worker.Messages <- msg
+
+	if cb != nil {
+		worker.Callbacks[msg] = cb
+	}
+}
+
+func (worker *Worker) ProcessMessage(msg WorkerMessage) WorkerMessage {
+
+	worker.Logger.Printf("<= %T\n", msg)
+	if cb, ok := worker.Callbacks[msg.InResponseTo()]; ok {
+		cb(msg)
+		delete(worker.Callbacks, msg)
+	}
+	return msg
+}
+
+func (worker *Worker) ProcessAction(msg WorkerMessage) WorkerMessage {
+
+	worker.Logger.Printf("<- %T\n", msg)
+	if cb, ok := worker.Callbacks[msg.InResponseTo()]; ok {
+		cb(msg)
+		delete(worker.Callbacks, msg)
+	}
+	return msg
+}
diff --git a/worker/worker.go b/worker/worker.go
index b665884..439ab64 100644
--- a/worker/worker.go
+++ b/worker/worker.go
@@ -9,22 +9,24 @@ import (
 	"net/url"
 )
 
-type Worker interface {
-	GetMessages() chan types.WorkerMessage
-	PostAction(types.WorkerMessage)
-	Run()
-}
-
 // Guesses the appropriate worker type based on the given source string
-func NewWorker(source string, logger *log.Logger) (Worker, error) {
+func NewWorker(source string, logger *log.Logger) (*types.Worker, error) {
 	u, err := url.Parse(source)
 	if err != nil {
 		return nil, err
 	}
+	worker := &types.Worker{
+		Actions:   make(chan types.WorkerMessage, 50),
+		Callbacks: make(map[types.WorkerMessage]func(msg types.WorkerMessage)),
+		Messages:  make(chan types.WorkerMessage, 50),
+		Logger:    logger,
+	}
 	switch u.Scheme {
 	case "imap":
 	case "imaps":
-		return imap.NewIMAPWorker(logger), nil
+		worker.Backend = imap.NewIMAPWorker(worker)
+	default:
+		return nil, fmt.Errorf("Unknown backend %s", u.Scheme)
 	}
-	return nil, fmt.Errorf("Unknown backend %s", u.Scheme)
+	return worker, nil
 }