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) 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 {
					writeResult(w, m, ldap.LDAPResultOperationsError, "internal server error")
				}
			}
		}()
		h(tw, m)
	}
}

// writeResult sends the result message that terminates a request,
// using the response type that pairs with the request type. It does
// nothing for operations that are never answered, such as abandon.
func writeResult(w ldap.ResponseWriter, m *ldap.Message, code int, diagnostic string) {
	if m == nil {
		return
	}
	res, ok := resultForRequest(m.ProtocolOp(), code, diagnostic)
	if !ok {
		return
	}
	w.Write(res)
}

// resultForRequest builds the result message that terminates a
// request. An LDAP client only accepts the response type that pairs
// with the request it sent (a search for instance is only completed
// by a searchResDone) and silently ignores anything else, so answering
// a search with a bare LDAPResult leaves the client waiting forever
// rather than reporting the error. The boolean reports whether a
// response is expected for the operation at all.
//
// Diagnostic messages are only attached where goldap exposes a setter
// for the type; the modify family carries the result code alone.
func resultForRequest(op message.ProtocolOp, code int, diagnostic string) (message.ProtocolOp, bool) {
	switch op.(type) {
	case message.SearchRequest:
		res := ldap.NewSearchResultDoneResponse(code)
		res.SetDiagnosticMessage(diagnostic)
		return res, true
	case message.BindRequest:
		res := ldap.NewBindResponse(code)
		res.SetDiagnosticMessage(diagnostic)
		return res, true
	case message.ExtendedRequest:
		res := ldap.NewExtendedResponse(code)
		res.SetDiagnosticMessage(diagnostic)
		return res, true
	case message.CompareRequest:
		return message.CompareResponse(ldap.NewResponse(code)), true
	case message.AddRequest:
		return message.AddResponse(ldap.NewResponse(code)), true
	case message.DelRequest:
		return message.DelResponse(ldap.NewResponse(code)), true
	case message.ModifyRequest:
		return message.ModifyResponse(ldap.NewResponse(code)), true
	case message.ModifyDNRequest:
		return message.ModifyDNResponse(ldap.NewResponse(code)), true
	}
	return nil, false
}

// 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)
}