// Copyright 2021 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:build linux
// +build linux

package sharedmem

import (
	"gvisor.dev/gvisor/pkg/atomicbitops"
	"gvisor.dev/gvisor/pkg/buffer"
	"gvisor.dev/gvisor/pkg/rawfile"
	"gvisor.dev/gvisor/pkg/sync"
	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/header"
	"gvisor.dev/gvisor/pkg/tcpip/stack"
)

type serverEndpoint struct {
	// bufferSize is the size of each individual buffer.
	// bufferSize is immutable.
	bufferSize uint32

	// rx is the receive queue.
	rx serverRx

	// stopRequested determines whether the worker goroutines should stop.
	stopRequested atomicbitops.Uint32

	// Wait group used to indicate that all workers have stopped.
	completed sync.WaitGroup

	// peerFD is an fd to the peer that can be used to detect when the peer is
	// gone.
	// peerFD is immutable.
	peerFD int

	// caps holds the endpoint capabilities.
	caps stack.LinkEndpointCapabilities

	// hdrSize is the size of the link layer header if any.
	// hdrSize is immutable.
	hdrSize uint32

	// virtioNetHeaderRequired if true indicates that a virtio header is expected
	// in all inbound/outbound packets.
	virtioNetHeaderRequired bool

	// onClosed is a function to be called when the FD's peer (if any) closes its
	// end of the communication pipe.
	onClosed func(tcpip.Error)

	// mu protects the following fields.
	mu sync.RWMutex

	// tx is the transmit queue.
	// +checklocks:mu
	tx serverTx

	// workerStarted specifies whether the worker goroutine was started.
	// +checklocks:mu
	workerStarted bool

	// addr is the local address of this endpoint.
	//
	// +checklocks:mu
	addr tcpip.LinkAddress
	// mtu (maximum transmission unit) is the maximum size of a packet.
	// +checklocks:mu
	mtu uint32
}

// NewServerEndpoint creates a new shared-memory-based endpoint. Buffers will be
// broken up into buffers of "bufferSize" bytes.
func NewServerEndpoint(opts Options) (stack.LinkEndpoint, error) {
	e := &serverEndpoint{
		mtu:        opts.MTU,
		bufferSize: opts.BufferSize,
		addr:       opts.LinkAddress,
		peerFD:     opts.PeerFD,
		onClosed:   opts.OnClosed,
	}

	if err := e.tx.init(&opts.RX); err != nil {
		return nil, err
	}

	if err := e.rx.init(&opts.TX); err != nil {
		e.tx.cleanup()
		return nil, err
	}

	e.caps = stack.LinkEndpointCapabilities(0)
	if opts.RXChecksumOffload {
		e.caps |= stack.CapabilityRXChecksumOffload
	}

	if opts.TXChecksumOffload {
		e.caps |= stack.CapabilityTXChecksumOffload
	}

	if opts.LinkAddress != "" {
		e.hdrSize = header.EthernetMinimumSize
		e.caps |= stack.CapabilityResolutionRequired
	}

	return e, nil
}

// SetOnCloseAction implements stack.LinkEndpoint.SetOnCloseAction.
func (*serverEndpoint) SetOnCloseAction(func()) {}

// Close frees all resources associated with the endpoint.
func (e *serverEndpoint) Close() {
	// Tell dispatch goroutine to stop, then write to the eventfd so that it wakes
	// up in case it's sleeping.
	e.stopRequested.Store(1)
	e.rx.eventFD.Notify()

	// Cleanup the queues inline if the worker hasn't started yet; we also know it
	// won't start from now on because stopRequested is set to 1.
	e.mu.Lock()
	defer e.mu.Unlock()
	workerPresent := e.workerStarted

	if !workerPresent {
		e.tx.cleanup()
		e.rx.cleanup()
	}
}

// Wait implements stack.LinkEndpoint.Wait. It waits until all workers have
// stopped after a Close() call.
func (e *serverEndpoint) Wait() {
	e.completed.Wait()
}

// Attach implements stack.LinkEndpoint.Attach. It launches the goroutine that
// reads packets from the rx queue.
func (e *serverEndpoint) Attach(dispatcher stack.NetworkDispatcher) {
	e.mu.Lock()
	if !e.workerStarted && e.stopRequested.Load() == 0 {
		e.workerStarted = true
		e.completed.Add(1)
		if e.peerFD >= 0 {
			e.completed.Add(1)
			// Spin up a goroutine to monitor for peer shutdown.
			go func() {
				b := make([]byte, 1)
				// When sharedmem endpoint is in use the peerFD is never used for any
				// data transfer and this Read should only return if the peer is
				// shutting down.
				_, errno := rawfile.BlockingRead(e.peerFD, b)
				if e.onClosed != nil {
					if errno == 0 {
						e.onClosed(nil)
					} else {
						e.onClosed(tcpip.TranslateErrno(errno))
					}
				}
				e.completed.Done()
			}()
		}
		// Link endpoints are not savable. When transportation endpoints are saved,
		// they stop sending outgoing packets and all incoming packets are rejected.
		go e.dispatchLoop(dispatcher) // S/R-SAFE: see above.
	}
	e.mu.Unlock()
}

// IsAttached implements stack.LinkEndpoint.IsAttached.
func (e *serverEndpoint) IsAttached() bool {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.workerStarted
}

// MTU implements stack.LinkEndpoint.MTU.
func (e *serverEndpoint) MTU() uint32 {
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.mtu
}

func (e *serverEndpoint) SetMTU(mtu uint32) {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.mtu = mtu
}

// Capabilities implements stack.LinkEndpoint.Capabilities.
func (e *serverEndpoint) Capabilities() stack.LinkEndpointCapabilities {
	return e.caps
}

// MaxHeaderLength implements stack.LinkEndpoint.MaxHeaderLength. It returns the
// ethernet frame header size.
func (e *serverEndpoint) MaxHeaderLength() uint16 {
	return uint16(e.hdrSize)
}

// LinkAddress implements stack.LinkEndpoint.LinkAddress. It returns the local
// link address.
func (e *serverEndpoint) LinkAddress() tcpip.LinkAddress {
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.addr
}

// SetLinkAddress implements stack.LinkEndpoint.SetLinkAddress.
func (e *serverEndpoint) SetLinkAddress(addr tcpip.LinkAddress) {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.addr = addr
}

// AddHeader implements stack.LinkEndpoint.AddHeader.
func (e *serverEndpoint) AddHeader(pkt *stack.PacketBuffer) {
	e.mu.RLock()
	defer e.mu.RUnlock()
	// Add ethernet header if needed.
	if len(e.addr) == 0 {
		return
	}

	eth := header.Ethernet(pkt.LinkHeader().Push(header.EthernetMinimumSize))
	eth.Encode(&header.EthernetFields{
		SrcAddr: pkt.EgressRoute.LocalLinkAddress,
		DstAddr: pkt.EgressRoute.RemoteLinkAddress,
		Type:    pkt.NetworkProtocolNumber,
	})
}

func (e *serverEndpoint) parseHeader(pkt *stack.PacketBuffer) bool {
	_, ok := pkt.LinkHeader().Consume(header.EthernetMinimumSize)
	return ok
}

// ParseHeader implements stack.LinkEndpoint.ParseHeader.
func (e *serverEndpoint) ParseHeader(pkt *stack.PacketBuffer) bool {
	e.mu.RLock()
	defer e.mu.RUnlock()
	// Add ethernet header if needed.
	if len(e.addr) == 0 {
		return true
	}

	return e.parseHeader(pkt)
}

func (e *serverEndpoint) AddVirtioNetHeader(pkt *stack.PacketBuffer) {
	virtio := header.VirtioNetHeader(pkt.VirtioNetHeader().Push(header.VirtioNetHeaderSize))
	virtio.Encode(&header.VirtioNetHeaderFields{})
}

// +checklocks:e.mu
func (e *serverEndpoint) writePacketLocked(r stack.RouteInfo, protocol tcpip.NetworkProtocolNumber, pkt *stack.PacketBuffer) tcpip.Error {
	if e.virtioNetHeaderRequired {
		e.AddVirtioNetHeader(pkt)
	}

	ok := e.tx.transmit(pkt)
	if !ok {
		return &tcpip.ErrWouldBlock{}
	}

	return nil
}

// WritePacket writes outbound packets to the file descriptor. If it is not
// currently writable, the packet is dropped.
// WritePacket implements stack.LinkEndpoint.WritePacket.
func (e *serverEndpoint) WritePacket(_ stack.RouteInfo, _ tcpip.NetworkProtocolNumber, pkt *stack.PacketBuffer) tcpip.Error {
	// Transmit the packet.
	e.mu.Lock()
	defer e.mu.Unlock()
	if err := e.writePacketLocked(pkt.EgressRoute, pkt.NetworkProtocolNumber, pkt); err != nil {
		return err
	}
	e.tx.notify()
	return nil
}

// WritePackets implements stack.LinkEndpoint.WritePackets.
func (e *serverEndpoint) WritePackets(pkts stack.PacketBufferList) (int, tcpip.Error) {
	n := 0
	var err tcpip.Error
	e.mu.Lock()
	defer e.mu.Unlock()
	for _, pkt := range pkts.AsSlice() {
		if err = e.writePacketLocked(pkt.EgressRoute, pkt.NetworkProtocolNumber, pkt); err != nil {
			break
		}
		n++
	}
	// WritePackets never returns an error if it successfully transmitted at least
	// one packet.
	if err != nil && n == 0 {
		return 0, err
	}
	e.tx.notify()
	return n, nil
}

// dispatchLoop reads packets from the rx queue in a loop and dispatches them
// to the network stack.
func (e *serverEndpoint) dispatchLoop(d stack.NetworkDispatcher) {
	for e.stopRequested.Load() == 0 {
		b := e.rx.receive()
		if b == nil {
			e.rx.EnableNotification()
			// Now pull again to make sure we didn't receive any packets
			// while notifications were not enabled.
			for {
				b = e.rx.receive()
				if b != nil {
					// Disable notifications as we only need to be notified when we are going
					// to block on eventFD. This should prevent the peer from needlessly
					// writing to eventFD when this end is already awake and processing
					// packets.
					e.rx.DisableNotification()
					break
				}
				e.rx.waitForPackets()
			}
		}
		pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
			Payload: buffer.MakeWithView(b),
		})
		if e.virtioNetHeaderRequired {
			_, ok := pkt.VirtioNetHeader().Consume(header.VirtioNetHeaderSize)
			if !ok {
				pkt.DecRef()
				continue
			}
		}
		var proto tcpip.NetworkProtocolNumber
		e.mu.RLock()
		addrLen := len(e.addr)
		e.mu.RUnlock()
		if addrLen != 0 {
			if !e.parseHeader(pkt) {
				pkt.DecRef()
				continue
			}
			proto = header.Ethernet(pkt.LinkHeader().Slice()).Type()
		} else {
			// We don't get any indication of what the packet is, so try to guess
			// if it's an IPv4 or IPv6 packet.
			// IP version information is at the first octet, so pulling up 1 byte.
			h, ok := pkt.Data().PullUp(1)
			if !ok {
				pkt.DecRef()
				continue
			}
			switch header.IPVersion(h) {
			case header.IPv4Version:
				proto = header.IPv4ProtocolNumber
			case header.IPv6Version:
				proto = header.IPv6ProtocolNumber
			default:
				pkt.DecRef()
				continue
			}
		}
		// Send packet up the stack.
		d.DeliverNetworkPacket(proto, pkt)
		pkt.DecRef()
	}

	e.mu.Lock()
	defer e.mu.Unlock()

	// Clean state.
	e.tx.cleanup()
	e.rx.cleanup()

	e.completed.Done()
}

// ARPHardwareType implements stack.LinkEndpoint.ARPHardwareType
func (e *serverEndpoint) ARPHardwareType() header.ARPHardwareType {
	if e.hdrSize > 0 {
		return header.ARPHardwareEther
	}
	return header.ARPHardwareNone
}