// -*- Mode: Go; indent-tabs-mode: t -*-

/*
 * Copyright (C) 2023-2025 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

// Package listener implements a high-level interface to the apparmor
// notification mechanism. It can be used to build userspace applications
// which respond to apparmor prompting profiles.
package listener

import (
	"errors"
	"fmt"
	"os"
	"sync"

	"github.com/snapcore/snapd/logger"
	"github.com/snapcore/snapd/osutil/epoll"
	"github.com/snapcore/snapd/sandbox/apparmor"
	"github.com/snapcore/snapd/sandbox/apparmor/notify"
)

var (
	// ErrClosed indicates that the listener has been closed.
	ErrClosed = errors.New("listener has been closed")

	// ErrAlreadyClosed indicates that the listener has previously been closed.
	ErrAlreadyClosed = errors.New("listener has already been closed")

	// ErrNotSupported indicates that the kernel does not support apparmor prompting.
	ErrNotSupported = errors.New("kernel does not support apparmor notifications")

	osOpen                       = os.Open
	notifyRegisterFileDescriptor = notify.RegisterFileDescriptor
	notifyIoctl                  = notify.Ioctl
)

// SendResponseFunc sends a response to the kernel on behalf of a request.
type SendResponseFunc = func(id uint64, aaAllowed, aaRequested, allow notify.AppArmorPermission) error

// Listener encapsulates a loop for receiving apparmor notification requests
// and responding with notification responses, hiding the low-level details.
type Listener[R any] struct {
	// runOnce ensures that the Run method is only run once, to avoid closing
	// a channel multiple times
	runOnce sync.Once

	// newRequest is a closure which the listener can call to construct a new
	// request when it receives a message from the kernel. The listener should
	// always call newRequest with l.buildAndSendResponse passed in as the
	// sendResponse closure.
	newRequest func(msg notify.MsgNotificationGeneric, sendResponse SendResponseFunc) (*R, error)

	// reqs is a channel over which to send requests to the manager.
	// Only the main run loop may close this channel.
	reqs chan *R

	// readyOnce should always be used to call `signalReady` to ensure that the
	// listener only readies once.
	readyOnce sync.Once
	// ready is a channel which is closed once all requests which were pending
	// at time of registration have been re-received from the kernel and sent
	// over the reqs channel. This occurs once pendingCount reaches 0 or snapd
	// receives a message which does not have the UNOTIF_RESENT flag. This
	// channel must only be closed by the signalReady method, and that should
	// only be called via readyOnce.
	ready chan struct{}
	// pendingMu is a mutex which protects pendingCount.
	pendingMu sync.Mutex
	// pendingCount is the number of "pending" (UNOTIF_RESENT) messages still
	// expected to be re-received from the kernel. When the listener becomes
	// ready, this count must be sent to 0.
	pendingCount int

	// protocolVersion is the notification protocol version associated with the
	// listener's notify socket. Once registered with a particular version,
	// that version will be used for all messages sent or received over that
	// socket.
	protocolVersion notify.ProtocolVersion

	// socketMu guards the notify file, the epoll instance, and the close chan.
	socketMu   sync.Mutex
	notifyFile *os.File
	poll       *epoll.Epoll
	// closeChan will be closed by Close to indicate to the run loop that the
	// listener should be closed.
	closeChan chan struct{}
}

// Register opens and configures the apparmor notification interface.
//
// If the kernel does not support the notification mechanism the error is ErrNotSupported.
func Register[R any](newReq func(msg notify.MsgNotificationGeneric, sendResponse SendResponseFunc) (*R, error)) (listener *Listener[R], err error) {
	path := apparmor.NotifySocketPath
	if override := os.Getenv("PROMPT_NOTIFY_PATH"); override != "" {
		path = override
	}

	notifyFile, err := osOpen(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, ErrNotSupported
		}
		return nil, fmt.Errorf("cannot open %q: %v", path, err)
	}
	defer func() {
		if err != nil {
			notifyFile.Close()
		}
	}()

	poll, err := epoll.Open()
	if err != nil {
		return nil, fmt.Errorf("cannot open epoll file descriptor: %v", err)
	}
	defer func() {
		if err != nil {
			poll.Close()
		}
	}()
	if err = poll.Register(int(notifyFile.Fd()), epoll.Readable); err != nil {
		return nil, fmt.Errorf("cannot register epoll on %q: %v", path, err)
	}

	protoVersion, pendingCount, err := notifyRegisterFileDescriptor(notifyFile.Fd())
	if err != nil {
		return nil, err
	}
	logger.Debugf("registered listener with protocol version %d", protoVersion)

	listener = &Listener[R]{
		newRequest: newReq,

		reqs: make(chan *R),

		ready:        make(chan struct{}),
		pendingCount: pendingCount,

		protocolVersion: protoVersion,

		notifyFile: notifyFile,
		poll:       poll,
		closeChan:  make(chan struct{}),
	}
	// If there are no pending requests waiting to be re-sent, ready
	// immediately, otherwise start the ready timer when Run is called.
	if listener.pendingCount == 0 {
		listener.readyOnce.Do(func() {
			listener.signalReady()
		})
	}
	return listener, nil
}

// isClosed returns true if the listener has been closed.
//
// Any caller which must ensure that a Close is not in progress should hold the
// lock while checking isClosed, and continue to hold the lock until it is safe
// for Close to be run.
func (l *Listener[R]) isClosed() bool {
	select {
	case <-l.closeChan:
		return true
	default:
		return false
	}
}

// Close stops the listener and closes the kernel communication file.
func (l *Listener[R]) Close() error {
	l.socketMu.Lock()
	defer l.socketMu.Unlock()
	if l.isClosed() {
		return ErrAlreadyClosed
	}

	// Closing the notify file signals to the kernel that the listener is
	// disconnecting, so the kernel will send back denials or pass requests
	// on to other listeners which connect. Do this before closing the epoll
	// instance so that the kernel does not try to send any further messages
	// which won't be received.
	err1 := l.notifyFile.Close()

	// Close the close channel so that the the run loop knows to stop trying
	// to send requests over the request channel, and so that once the epoll
	// FD is closed (causing the syscall to error), it can check the closeChan
	// to see whether Close was called or whether a real error occurred.
	close(l.closeChan)

	// Close the epoll so that if the run loop is waiting on an event, it will
	// return an error.
	err2 := listenerEpollClose(l)
	if err1 != nil {
		return err1
	}
	return err2
}

// Reqs returns a read-only channel through which requests may be received.
// The channel is closed when the Close() method is called or an error occurs.
func (l *Listener[R]) Reqs() <-chan *R {
	return l.reqs
}

// Ready returns a read-only channel which will be closed once all requests
// which were pending when the listener was registered have been re-received
// from the kernel and sent over the reqs channel.
//
// The kernel guarantees that no non-resent requests will be sent until all
// originally-pending requests have been resent, so if a non-resent request is
// received, then this channel will be closed to signal readiness.
func (l *Listener[R]) Ready() <-chan struct{} {
	return l.ready
}

// Allow tests to kill the listener instead of logging errors so that tests
// don't have race condition failures.
var exitOnError = false

// Run reads and dispatches kernel requests until the listener is closed.
//
// Run should only be called once per listener object, and it runs until the
// listener is closed or errors (if exitOnError is true), and returns the cause.
// If the listener was intentionally stopped via the Close() method, returns nil.
func (l *Listener[R]) Run() error {
	var err error
	l.runOnce.Do(func() {
		// Run should only be called once, so this runOnce.Do is really only an
		// extra precaution to ensure that l.reqs is only closed once.
		defer func() {
			// When listener run loop ends, close the requests channel.
			close(l.reqs)
			// The manager is closing the listener, so if the listener isn't
			// ready by now, it doesn't matter, the manager has stopped
			// receiving requests by now anyway.
		}()
		for {
			err = l.handleRequests()
			if err != nil {
				if errors.Is(err, ErrClosed) {
					// Don't treat the listener closing as a real error
					err = nil
					return
				} else if exitOnError {
					l.Close() // make sure Close is called at least once
					return
				}
				logger.Noticef("error in prompting listener run loop: %v", err)
			}
		}
	})
	return err
}

// epollWaiter allows listenerEpollWait to be mocked in tests.
type epollWaiter interface {
	epollWait() ([]epoll.Event, error)
	epollClose() error
	socketFD() int
}

var listenerEpollWait = func(l epollWaiter) ([]epoll.Event, error) {
	return l.epollWait()
}

func (l *Listener[R]) epollWait() ([]epoll.Event, error) {
	return l.poll.Wait()
}

var listenerEpollClose = func(l epollWaiter) error {
	return l.epollClose()
}

func (l *Listener[R]) epollClose() error {
	return l.poll.Close()
}

func (l *Listener[R]) socketFD() int {
	l.socketMu.Lock()
	defer l.socketMu.Unlock()
	return int(l.notifyFile.Fd())
}

func (l *Listener[R]) handleRequests() error {
	// Get the socket FD with the lock held, so we don't break our contract.
	// Do this before any check whether the epoll instance is closed, so we
	// are sure to have the pre-close FD in case the instance is closed after
	// the check.
	socketFd := l.socketFD()

	events, err := listenerEpollWait(l)
	if err != nil {
		// The epoll syscall returned an error, so let's see whether it was
		// because we closed the epoll FD.
		if l.isClosed() {
			return ErrClosed
		}
		return err
	}
	// XXX: Since queued notifications should be sent immediately by the kernel,
	// if we think there are pending requests, we could use a non-blocking read
	// to poll for new requests, and if we don't immediately see one, we could
	// use a (not yet implemented) ioctl command to check the listener's status
	// regarding ready/pending requests and update the listener's pending count.
	// This would require some work on the kernel side, so it could be a future
	// enhancement, but not one we can pursue at time of writing.

	for _, event := range events {
		if event.Fd != socketFd {
			logger.Debugf("unexpected event from fd %v (%v)", event.Fd, event.Readiness)
			continue
		}
		if event.Readiness&epoll.Readable == 0 {
			continue
		}
		// Prepare a receive buffer for incoming request. The buffer is of the
		// maximum allowed size and will contain one or more kernel requests
		// upon return.
		ioctlBuf := notify.NewIoctlRequestBuffer(l.protocolVersion)
		buf, err := l.doIoctl(notify.APPARMOR_NOTIF_RECV, ioctlBuf)
		if err != nil {
			return err
		}
		if err := l.decodeAndDispatchRequest(buf); err != nil {
			return err
		}
	}
	return nil
}

// doIoctl locks the mutex guarding the notify socket, checks whether the
// listener is being closed, and if not, sends an ioctl request with the given
// request type and buffer.
func (l *Listener[R]) doIoctl(sendOrRecv notify.IoctlRequest, buf notify.IoctlRequestBuffer) ([]byte, error) {
	l.socketMu.Lock()
	defer l.socketMu.Unlock()
	if l.isClosed() {
		return nil, ErrClosed
	}
	return notifyIoctl(l.notifyFile.Fd(), sendOrRecv, buf)
}

// decodeAndDispatchRequest reads all messages from the given buffer, decodes
// each one into a message notification for the particular mediation class,
// creates a Request from that message, and attempts to send it to the manager
// via the request channel.
func (l *Listener[R]) decodeAndDispatchRequest(buf []byte) error {
	for len(buf) > 0 {
		first, rest, err := notify.ExtractFirstMsg(buf)
		if err != nil {
			return err
		}
		buf = rest

		var nmsg notify.MsgNotification
		if err := nmsg.UnmarshalBinary(first); err != nil {
			return err
		}
		if nmsg.Version != l.protocolVersion {
			return fmt.Errorf("unexpected protocol version: listener registered with %d, but received %d", l.protocolVersion, nmsg.Version)
		}
		// What kind of notification message did we get? (I hope it's an Op)
		if nmsg.NotificationType != notify.APPARMOR_NOTIF_OP {
			return fmt.Errorf("unsupported notification type: %v", nmsg.NotificationType)
		}
		var omsg notify.MsgNotificationOp
		if err := omsg.UnmarshalBinary(first); err != nil {
			return err
		}

		// Handle whether the message was resent prior to attempting to parse
		// it into a Request or send it to the manager. This way, API calls
		// won't block unnecessarily. If this is the final pending request,
		// then we want to signal readiness only if we find it to be malformed
		// or after the request has been received, to make sure the manager
		// handles the request (if valid) prior to observing the listener
		// readiness.
		isFinalPendingReq := l.handlePotentialResentMessage(omsg.Resent())

		req, err := l.parseRequest(omsg.Class, first)
		if err != nil {
			// Auto-deny it now, else the triggering application will just hang
			// until the message times out in the kernel.
			l.denyMalformedRequest(&omsg)
			if isFinalPendingReq {
				// In practice, snapd should always be able to parse a message
				// with a RESENT flag set unless the list of classes/interfaces/permissions
				// supported by snapd changed across a restart.
				l.readyOnce.Do(func() {
					l.signalReady()
				})
			}
			return err
		}

		select {
		case l.reqs <- req:
			// request received
		case <-l.closeChan:
			// The listener is being closed, so stop trying to deliver the
			// message up to the manager. It will appear to the kernel that we
			// have received and processed the request, when in reality, we
			// have not. The higher-level restart handling logic ensures that
			// the request will be re-sent to snapd when it restarts and
			// re-registers the listener.
			return ErrClosed
		}

		if isFinalPendingReq {
			l.readyOnce.Do(func() {
				l.signalReady()
			})
		}
	}
	return nil
}

func (l *Listener[R]) parseRequest(class notify.MediationClass, buf []byte) (*R, error) {
	var err error
	var msg notify.MsgNotificationGeneric
	switch class {
	case notify.AA_CLASS_FILE:
		msg, err = parseMsgNotificationFile(buf)
	default:
		err = fmt.Errorf("unsupported mediation class: %v", class)
	}
	if err != nil {
		return nil, err
	}

	return l.newRequest(msg, l.buildAndSendResponse)
}

func parseMsgNotificationFile(buf []byte) (*notify.MsgNotificationFile, error) {
	var fmsg notify.MsgNotificationFile
	if err := fmsg.UnmarshalBinary(buf); err != nil {
		return nil, err
	}
	logger.Debugf("received file request from the kernel: %+v", fmsg)
	return &fmsg, nil
}

func (l *Listener[R]) denyMalformedRequest(msg *notify.MsgNotificationOp) {
	resp := msg.BuildDenyResponse()
	l.encodeAndSendResponse(resp)
}

// handlePotentialResentMessage handles whether the message has been received
// with the RESENT flag set or not. If so, then decrement the pending count and
// return true if this is the final pending request. If not resent, then ensure
// that the listener has signalled readiness.
//
// The kernel is guaranteed to resend all previously-pending messages before it
// sends any new messages, so if it's NOT resent, then we know the kernel is
// done sending pending messages, and we should ready immediately.
func (l *Listener[R]) handlePotentialResentMessage(resent bool) (isFinalPendingReq bool) {
	if resent {
		return l.decrementPendingCheckFinal()
	}
	l.readyOnce.Do(func() {
		if pendingCount := l.signalReady(); pendingCount != 0 {
			logger.Noticef("received non-resent message when pending count was %d", pendingCount)
		}
	})
	return false
}

// decrementPendingCheckFinal decrements the pending count if it's not already
// 0, and returns whether this was the final pending request.
//
// The caller should only call this method if the message associated with the
// request *was* marked by the kernel as having been previously sent.
func (l *Listener[R]) decrementPendingCheckFinal() (isFinal bool) {
	l.pendingMu.Lock()
	defer l.pendingMu.Unlock()
	if l.pendingCount == 0 {
		return false
	}
	l.pendingCount--
	if l.pendingCount == 0 {
		return true
	}
	return false
}

// signalReady is responsible for closing the ready channel and ensuring that
// pendingCount is set to 0. Returns the pending count at time of readying.
//
// Potential callers must ensure that this method is only called once per
// listener.
func (l *Listener[R]) signalReady() (pendingCount int) {
	l.pendingMu.Lock()
	defer l.pendingMu.Unlock()
	pendingCount = l.pendingCount
	l.pendingCount = 0 // tell the run loop we're ready
	close(l.ready)
	return pendingCount
}

func (l *Listener[R]) buildAndSendResponse(id uint64, aaAllowed, aaRequested, allow notify.AppArmorPermission) error {
	resp := notify.BuildResponse(l.protocolVersion, id, aaAllowed, aaRequested, allow)
	return l.encodeAndSendResponse(resp)
}

func (l *Listener[R]) encodeAndSendResponse(resp *notify.MsgNotificationResponse) error {
	buf, err := resp.MarshalBinary()
	if err != nil {
		return err
	}
	ioctlBuf := notify.IoctlRequestBuffer(buf)
	_, err = l.doIoctl(notify.APPARMOR_NOTIF_SEND, ioctlBuf)
	return err
}
