/*
 * Copyright (C) 2022 The Android Open Source Project
 *
 * 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.
 */

#pragma once

#include <android-base/result.h>
#include <android-base/unique_fd.h>
#include <linux/bpf.h>
#include <sys/mman.h>
#include <utils/Log.h>

#include "bpf/BpfUtils.h"

namespace android {
namespace bpf {

// BpfRingbufBase contains the non-templated functionality of BPF ring buffers.
class BpfRingbufBase {
 public:
  ~BpfRingbufBase() {
    if (mConsumerPos) munmap(mConsumerPos, mConsumerSize);
    if (mProducerPos) munmap(mProducerPos, mProducerSize);
    mConsumerPos = nullptr;
    mProducerPos = nullptr;
  }

 protected:
  // Non-initializing constructor, used by Create.
  BpfRingbufBase(size_t value_size) : mValueSize(value_size) {}

  // Full construction that aborts on error (use Create/Init to handle errors).
  BpfRingbufBase(const char* path, size_t value_size) : mValueSize(value_size) {
    if (auto status = Init(path); !status.ok()) {
      ALOGE("BpfRingbuf init failed: %s", status.error().message().c_str());
      abort();
    }
  }

  // Delete copy constructor (class owns raw pointers).
  BpfRingbufBase(const BpfRingbufBase&) = delete;

  // Initialize the base ringbuffer components. Must be called exactly once.
  base::Result<void> Init(const char* path);

  // Consumes all messages from the ring buffer, passing them to the callback.
  base::Result<int> ConsumeAll(
      const std::function<void(const void*)>& callback);

  // Replicates c-style void* "byte-wise" pointer addition.
  template <typename Ptr>
  static Ptr pointerAddBytes(void* base, ssize_t offset_bytes) {
    return reinterpret_cast<Ptr>(reinterpret_cast<char*>(base) + offset_bytes);
  }

  // Rounds len by clearing bitmask, adding header, and aligning to 8 bytes.
  static uint32_t roundLength(uint32_t len) {
    len &= ~(BPF_RINGBUF_BUSY_BIT | BPF_RINGBUF_DISCARD_BIT);
    len += BPF_RINGBUF_HDR_SZ;
    return (len + 7) & ~7;
  }

  const size_t mValueSize;

  size_t mConsumerSize;
  size_t mProducerSize;
  unsigned long mPosMask;
  android::base::unique_fd mRingFd;

  void* mDataPos = nullptr;
  unsigned long* mConsumerPos = nullptr;
  unsigned long* mProducerPos = nullptr;
};

// This is a class wrapper for eBPF ring buffers. An eBPF ring buffer is a
// special type of eBPF map used for sending messages from eBPF to userspace.
// The implementation relies on fast shared memory and atomics for the producer
// and consumer management. Ring buffers are a faster alternative to eBPF perf
// buffers.
//
// This class is thread compatible, but not thread safe.
//
// Note: A kernel eBPF ring buffer may be accessed by both kernel and userspace
// processes at the same time. However, the userspace consumers of a given ring
// buffer all share a single read pointer. There is no guarantee which readers
// will read which messages.
template <typename Value>
class BpfRingbuf : public BpfRingbufBase {
 public:
  using MessageCallback = std::function<void(const Value&)>;

  // Creates a ringbuffer wrapper from a pinned path. This initialization will
  // abort on error. To handle errors, initialize with Create instead.
  BpfRingbuf(const char* path) : BpfRingbufBase(path, sizeof(Value)) {}

  // Creates a ringbuffer wrapper from a pinned path. There are no guarantees
  // that the ringbuf outputs messaged of type `Value`, only that they are the
  // same size. Size is only checked in ConsumeAll.
  static base::Result<std::unique_ptr<BpfRingbuf<Value>>> Create(
      const char* path);

  // Consumes all messages from the ring buffer, passing them to the callback.
  // Returns the number of messages consumed or a non-ok result on error. If the
  // ring buffer has no pending messages an OK result with count 0 is returned.
  base::Result<int> ConsumeAll(const MessageCallback& callback);

 private:
  // Empty ctor for use by Create.
  BpfRingbuf() : BpfRingbufBase(sizeof(Value)) {}
};

#define ACCESS_ONCE(x) (*(volatile typeof(x)*)&(x))

#if defined(__i386__) || defined(__x86_64__)
#define smp_sync() asm volatile("" ::: "memory")
#elif defined(__aarch64__)
#define smp_sync() asm volatile("dmb ish" ::: "memory")
#else
#define smp_sync() __sync_synchronize()
#endif

#define smp_store_release(p, v) \
  do {                          \
    smp_sync();                 \
    ACCESS_ONCE(*(p)) = (v);    \
  } while (0)

#define smp_load_acquire(p)        \
  ({                               \
    auto ___p = ACCESS_ONCE(*(p)); \
    smp_sync();                    \
    ___p;                          \
  })

inline base::Result<void> BpfRingbufBase::Init(const char* path) {
  if (sizeof(unsigned long) != 8) {
    return android::base::Error()
           << "BpfRingbuf does not support 32 bit architectures";
  }
  mRingFd.reset(mapRetrieveRW(path));
  if (!mRingFd.ok()) {
    return android::base::ErrnoError()
           << "failed to retrieve ringbuffer at " << path;
  }

  int map_type = android::bpf::bpfGetFdMapType(mRingFd);
  if (map_type != BPF_MAP_TYPE_RINGBUF) {
    errno = EINVAL;
    return android::base::ErrnoError()
           << "bpf map has wrong type: want BPF_MAP_TYPE_RINGBUF ("
           << BPF_MAP_TYPE_RINGBUF << ") got " << map_type;
  }

  int max_entries = android::bpf::bpfGetFdMaxEntries(mRingFd);
  if (max_entries < 0) {
    return android::base::ErrnoError()
           << "failed to read max_entries from ringbuf";
  }
  if (max_entries == 0) {
    errno = EINVAL;
    return android::base::ErrnoError() << "max_entries must be non-zero";
  }

  mPosMask = max_entries - 1;
  mConsumerSize = getpagesize();
  mProducerSize = getpagesize() + 2 * max_entries;

  {
    void* ptr = mmap(NULL, mConsumerSize, PROT_READ | PROT_WRITE, MAP_SHARED,
                     mRingFd, 0);
    if (ptr == MAP_FAILED) {
      return android::base::ErrnoError()
             << "failed to mmap ringbuf consumer pages";
    }
    mConsumerPos = reinterpret_cast<unsigned long*>(ptr);
  }

  {
    void* ptr = mmap(NULL, mProducerSize, PROT_READ, MAP_SHARED, mRingFd,
                     mConsumerSize);
    if (ptr == MAP_FAILED) {
      return android::base::ErrnoError()
             << "failed to mmap ringbuf producer page";
    }
    mProducerPos = reinterpret_cast<unsigned long*>(ptr);
  }

  mDataPos = pointerAddBytes<void*>(mProducerPos, getpagesize());
  return {};
}

inline base::Result<int> BpfRingbufBase::ConsumeAll(
    const std::function<void(const void*)>& callback) {
  int64_t count = 0;
  unsigned long cons_pos = smp_load_acquire(mConsumerPos);
  unsigned long prod_pos = smp_load_acquire(mProducerPos);
  while (cons_pos < prod_pos) {
    // Find the start of the entry for this read (wrapping is done here).
    void* start_ptr = pointerAddBytes<void*>(mDataPos, cons_pos & mPosMask);

    // The entry has an 8 byte header containing the sample length.
    uint32_t length = smp_load_acquire(reinterpret_cast<uint32_t*>(start_ptr));

    // If the sample isn't committed, we're caught up with the producer.
    if (length & BPF_RINGBUF_BUSY_BIT) return count;

    cons_pos += roundLength(length);

    if ((length & BPF_RINGBUF_DISCARD_BIT) == 0) {
      if (length != mValueSize) {
        smp_store_release(mConsumerPos, cons_pos);
        errno = EMSGSIZE;
        return android::base::ErrnoError()
               << "BPF ring buffer message has unexpected size (want "
               << mValueSize << " bytes, got " << length << " bytes)";
      }
      callback(pointerAddBytes<const void*>(start_ptr, BPF_RINGBUF_HDR_SZ));
      count++;
    }

    smp_store_release(mConsumerPos, cons_pos);
  }

  return count;
}

template <typename Value>
inline base::Result<std::unique_ptr<BpfRingbuf<Value>>>
BpfRingbuf<Value>::Create(const char* path) {
  auto rb = std::unique_ptr<BpfRingbuf>(new BpfRingbuf);
  if (auto status = rb->Init(path); !status.ok()) return status.error();
  return rb;
}

template <typename Value>
inline base::Result<int> BpfRingbuf<Value>::ConsumeAll(
    const MessageCallback& callback) {
  return BpfRingbufBase::ConsumeAll([&](const void* value) {
    callback(*reinterpret_cast<const Value*>(value));
  });
}

#undef ACCESS_ONCE
#undef smp_sync
#undef smp_store_release
#undef smp_load_acquire

}  // namespace bpf
}  // namespace android