package ldap

import (
	"github.com/vjeantet/goldap/message"
	ldap "github.com/vjeantet/ldapserver"
)

// guard wraps a handler so that a panic in handler code is logged and
// converted into an error response instead of killing the process. The
// underlying server runs each request in its own goroutine with no
// recovery of its own, so a single panic anywhere in a handler takes
// the whole daemon down with it.
//
// Requests that already wrote something cannot be recovered cleanly,
// since the client is expecting whatever it received to be terminated
// by the matching done message, so those are only logged.
func (s *server) guard(h ldap.HandlerFunc, fail func(ldap.ResponseWriter)) ldap.HandlerFunc {
	return func(w ldap.ResponseWriter, m *ldap.Message) {
		tw := &trackingWriter{ResponseWriter: w}
		defer func() {
			if r := recover(); r != nil {
				s.l.Error("Recovered from panic in handler", "panic", r)
				if !tw.written && fail != nil {
					fail(tw.ResponseWriter)
				}
			}
		}()
		h(tw, m)
	}
}

// failBindError, failSearchError and failOperationError write the
// terminating error response for each kind of operation and are used
// by guard when a handler dies before it started its reply.
func failBindError(w ldap.ResponseWriter) {
	res := ldap.NewBindResponse(ldap.LDAPResultOperationsError)
	res.SetDiagnosticMessage("internal server error")
	w.Write(res)
}

func failSearchError(w ldap.ResponseWriter) {
	res := ldap.NewSearchResultDoneResponse(ldap.LDAPResultOperationsError)
	res.SetDiagnosticMessage("internal server error")
	w.Write(res)
}

func failOperationError(w ldap.ResponseWriter) {
	res := ldap.NewResponse(ldap.LDAPResultOperationsError)
	res.SetDiagnosticMessage("internal server error")
	w.Write(res)
}

// trackingWriter records whether anything has been written so that
// guard can decide if an error response is still appropriate.
type trackingWriter struct {
	ldap.ResponseWriter
	written bool
}

func (w *trackingWriter) Write(po message.ProtocolOp) {
	w.written = true
	w.ResponseWriter.Write(po)
}