mayu-live 0.0.0

data/lib/mayu/message_cipher.rb

@@ -0,0 +1,172 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "sorbet-runtime"
+require "time"
+require "digest/sha2"
+require "openssl"
+require "securerandom"
+require "brotli"
+
+module Mayu
+  class MessageCipher
+    extend T::Sig
+
+    DEFAULT_TTL_SECONDS = T.let(10, Integer)
+
+    Message = T.type_alias { { iss: Float, exp: Float, payload: T.untyped } }
+
+    class Error < StandardError
+    end
+    class ExpiredError < Error
+    end
+    class IssuedInTheFutureError < Error
+    end
+    class EncryptError < Error
+    end
+    class DecryptError < Error
+    end
+    class InvalidHMACError < Error
+    end
+
+    sig { params(key: String, ttl: Integer).void }
+    def initialize(key:, ttl: DEFAULT_TTL_SECONDS)
+      raise ArgumentError, "ttl must be positive" unless ttl.positive?
+      @default_ttl_seconds = ttl
+      @key = T.let(Digest::SHA256.digest(key), String)
+    end
+
+    sig do
+      params(payload: T.untyped, auth_data: String, ttl: Integer).returns(
+        String
+      )
+    end
+    def dump(payload, auth_data: "", ttl: @default_ttl_seconds)
+      raise ArgumentError, "ttl must be positive" unless ttl.positive?
+      now = Time.now.to_f
+      message = { iss: now, exp: now + ttl, payload: Marshal.dump(payload) }
+      encode_message(message, auth_data:)
+    end
+
+    sig { params(data: String, auth_data: String).returns(T.untyped) }
+    def load(data, auth_data: "")
+      Marshal.load(decode_message(data, auth_data:))
+    end
+
+    private
+
+    sig { params(message: Message, auth_data: String).returns(String) }
+    def encode_message(message, auth_data: "")
+      message
+        .then { Marshal.dump(_1) }
+        .then { prepend_hmac(_1) }
+        .then { Brotli.deflate(_1) }
+        .then { encrypt(_1, auth_data:) }
+    end
+
+    sig { params(message: String, auth_data: String).returns(String) }
+    def decode_message(message, auth_data: "")
+      message
+        .then { decrypt(_1, auth_data:) }
+        .then { Brotli.inflate(_1) }
+        .then { validate_hmac(_1) }
+        .then { Marshal.load(_1) }
+        .tap { validate_times(_1) }
+        .fetch(:payload)
+    end
+
+    sig { params(input: String).returns(String) }
+    def prepend_hmac(input)
+      hmac = Digest::SHA256.digest(input)
+      input.prepend(hmac)
+    end
+
+    sig { params(input: String).returns(String) }
+    def validate_hmac(input)
+      hmac, message = input.unpack("a32 a*")
+
+      unless OpenSSL.fixed_length_secure_compare(
+               hmac,
+               Digest::SHA256.digest(message.to_s)
+             )
+        raise InvalidHMACError
+      end
+
+      message.to_s
+    end
+
+    sig { params(message: { iss: Float, exp: Float, payload: String }).void }
+    def validate_times(message)
+      message => { iss:, exp: }
+      now = Time.now.to_f
+      validate_iss(now, iss)
+      validate_exp(now, exp)
+    end
+
+    sig { params(now: Float, iss: Float).void }
+    def validate_iss(now, iss)
+      return if iss < now
+
+      raise IssuedInTheFutureError,
+            "The message was issued at #{Time.at(iss).iso8601}, which is in the future"
+    end
+
+    sig { params(now: Float, exp: Float).void }
+    def validate_exp(now, exp)
+      return if exp > now
+
+      raise ExpiredError,
+            "The message expired at #{Time.at(exp).iso8601}, which is in the past"
+    end
+
+    sig { params(message: String, auth_data: String).returns(String) }
+    def encrypt(message, auth_data: "")
+      cipher = OpenSSL::Cipher.new("aes-256-gcm")
+      cipher.encrypt
+      salt = SecureRandom.random_bytes(8)
+      cipher.key = generate_key(salt)
+      cipher.iv = iv = cipher.random_iv
+      cipher.auth_data = auth_data
+      cipher_text = cipher.update(message) + cipher.final
+      auth_tag = cipher.auth_tag
+      [auth_tag.bytesize, auth_tag, salt, iv, cipher_text].pack("C a* a* a* a*")
+    rescue OpenSSL::Cipher::CipherError
+      raise EncryptError
+    end
+
+    sig { params(data: String, auth_data: String).returns(String) }
+    def decrypt(data, auth_data: "")
+      data.unpack("C a*") => [Integer => auth_tag_len, String => data]
+
+      data.unpack("a#{auth_tag_len} a8 a12 a*") => [
+        auth_tag,
+        salt,
+        iv,
+        cipher_text
+      ]
+
+      cipher = OpenSSL::Cipher.new("aes-256-gcm")
+      cipher.iv = iv
+      cipher.key = generate_key(salt)
+      cipher.auth_data = auth_data
+      cipher.auth_tag = auth_tag
+      cipher.update(cipher_text) + cipher.final
+    rescue NoMatchingPatternError
+      raise DecryptError
+    rescue OpenSSL::Cipher::CipherError
+      raise DecryptError
+    end
+
+    sig { params(salt: String).returns(String) }
+    def generate_key(salt)
+      OpenSSL::KDF.scrypt(
+        @key,
+        salt:, # Salt.
+        N: 2**14, # CPU/memory cost parameter. This must be a power of 2.
+        r: 8, # Block size parameter.
+        p: 1, # Parallelization parameter
+        length: 32 # Length in octets of the derived key
+      )
+    end
+  end
+end

data/lib/mayu/message_cipher.test.rb

@@ -0,0 +1,16 @@
+# typed: true
+
+require "minitest/autorun"
+require "test_helper"
+
+require_relative "message_cipher"
+
+class TestMessageCipher < Minitest::Test
+  def test_dump_and_load
+    message_cipher = Mayu::MessageCipher.new(key: "test")
+
+    dumped = message_cipher.dump("hello")
+    loaded = message_cipher.load(dumped)
+    assert(loaded == "hello")
+  end
+end

data/lib/mayu/metrics/collector.rb

@@ -0,0 +1,161 @@
+# typed: strict
+
+module Mayu
+  module Metrics
+    module Collector
+      class DataStore
+        class MetricStore
+          extend T::Sig
+
+          sig { returns(DataStore) }
+          attr_reader :store
+
+          sig do
+            params(
+              store: DataStore,
+              metric_name: Symbol,
+              metric_type: Symbol,
+              metric_settings: T::Hash[Symbol, T.untyped]
+            ).void
+          end
+          def initialize(store, metric_name, metric_type:, metric_settings: {})
+            @store = store
+            @metric_name = metric_name
+            @metric_type = metric_type
+            @metric_settings = metric_settings
+            @aggregation_mode =
+              T.let(metric_settings.fetch(:aggregation, :sum), Symbol)
+          end
+
+          sig do
+            params(
+              val: T.any(Integer, Float),
+              labels: T::Hash[Symbol, T.untyped]
+            ).void
+          end
+          def set(val:, labels: {})
+          end
+
+          sig { returns(T::Hash[T::Array[Symbol], Float]) }
+          def all_values
+            @store
+              .values_for_metric(@metric_name)
+              .transform_values { aggregate(_1) }
+          end
+
+          private
+
+          sig { params(values: T::Array[Float]).returns(Float) }
+          def aggregate(values)
+            case @aggregation_mode
+            when :min
+              values.min
+            when :max
+              values.max
+            when :sum
+              values.sum
+            else
+              raise "Invalid aggregation setting"
+            end.to_f
+          end
+        end
+
+        extend T::Sig
+
+        sig { returns(InternalStore) }
+        attr_reader :internal_store
+
+        sig { params(internal_store: InternalStore).void }
+        def initialize(internal_store = {})
+          @internal_store = internal_store
+        end
+
+        sig do
+          params(metric_name: Symbol).returns(
+            T::Hash[T::Array[Symbol], T::Array[Float]]
+          )
+        end
+        def values_for_metric(metric_name)
+          @internal_store
+            .values
+            .map { _1[metric_name] }
+            .compact
+            .each_with_object(Hash.new { |h, k| h[k] = [] }) do |entries, obj|
+              entries.each { |labels, value| obj[labels] << value }
+            end
+        end
+
+        sig do
+          params(
+            metric_name: Symbol,
+            metric_type: Symbol,
+            metric_settings: T::Hash[Symbol, T.untyped]
+          ).returns(MetricStore)
+        end
+        def for_metric(metric_name, metric_type:, metric_settings: {})
+          MetricStore.new(self, metric_name, metric_type:, metric_settings: {})
+        end
+      end
+
+      class Server
+        extend T::Sig
+
+        sig { returns(Async::IO::Endpoint) }
+        attr_reader :endpoint
+
+        sig { params(endpoint: Async::IO::UNIXEndpoint).void }
+        def initialize(endpoint)
+          @endpoint = endpoint
+        end
+
+        sig { params(metric_name: Symbol).void }
+        def all_values(metric_name)
+          raise NotImplementedError, "Should this even be implemented?"
+        end
+
+        sig { void }
+        def start
+        end
+
+        sig { void }
+        def stop
+        end
+
+        sig do
+          params(
+            internal_store: InternalStore,
+            name: String,
+            restart: T::Boolean
+          ).void
+        end
+        def run(internal_store, name: self.class.name.to_s, restart: true)
+          wrapper = Wrapper.new
+
+          Console.logger.info(
+            self,
+            "Starting metrics collection on #{File.expand_path(@endpoint.path)}"
+          )
+
+          @endpoint.accept do |peer|
+            store = internal_store.store(peer, {})
+
+            unpacker = wrapper.unpacker(peer)
+
+            unpacker.each do |message|
+              case message
+              in [:store, data]
+                store.merge!(data)
+              else
+                Console
+                  .logger
+                  .warn(self) { "Unhandled mesage: #{message.inspect}" }
+              end
+            end
+          ensure
+            internal_store.delete(peer)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/mayu/metrics/exporter.rb

@@ -0,0 +1,47 @@
+# typed: strict
+
+module Mayu
+  module Metrics
+    class Exporter
+      class Server
+        extend T::Sig
+
+        sig do
+          params(
+            endpoint: Async::HTTP::Endpoint,
+            registry: Prometheus::Client::Registry
+          ).returns(Async::HTTP::Server)
+        end
+        def self.setup(endpoint:, registry:)
+          Console.logger.info(
+            self,
+            "Starting metrics exporter on #{endpoint.to_url}"
+          )
+
+          Async::HTTP::Server.for(
+            endpoint,
+            protocol: Async::HTTP::Protocol::HTTP11
+          ) do |request|
+            if request.path == "/favicon.ico"
+              next(
+                Protocol::HTTP::Response[
+                  404,
+                  { "content-type": "text/plain" },
+                  ["Not found"]
+                ]
+              )
+            end
+
+            body = Prometheus::Client::Formats::Text.marshal(registry)
+
+            Protocol::HTTP::Response[
+              200,
+              { "content-type": "text/plain" },
+              [body]
+            ]
+          end
+        end
+      end
+    end
+  end
+end

data/lib/mayu/metrics/reporter.rb

@@ -0,0 +1,187 @@
+# typed: strict
+
+module Mayu
+  module Metrics
+    module Reporter
+      extend T::Sig
+
+      sig do
+        type_parameters(:M)
+          .params(
+            collector_endpoint: Async::IO::UNIXEndpoint,
+            block:
+              T
+                .proc
+                .params(arg0: Prometheus::Client::Registry)
+                .returns(T.type_parameter(:M))
+          )
+          .returns(T.type_parameter(:M))
+      end
+      def self.run(collector_endpoint, &block)
+        data_store = DataStore.new
+        Prometheus::Client.config.data_store = data_store
+        metrics = yield(Prometheus::Client::Registry.new)
+        Client.connect_and_sync(collector_endpoint:, data_store:, interval: 1)
+        metrics
+      end
+
+      class Client
+        extend T::Sig
+
+        sig do
+          params(
+            collector_endpoint: Async::IO::UNIXEndpoint,
+            block: T.proc.params(arg0: Client).void
+          ).void
+        end
+        def self.connect(collector_endpoint, &block)
+          Console.logger.info(
+            self,
+            "Connecting to #{File.expand_path(collector_endpoint.path)}"
+          )
+
+          collector_endpoint.connect { |peer| yield new(peer) }
+        end
+
+        sig do
+          params(
+            collector_endpoint: Async::IO::UNIXEndpoint,
+            data_store: DataStore,
+            interval: Integer,
+            task: Async::Task
+          ).returns(Async::Task)
+        end
+        def self.connect_and_sync(
+          collector_endpoint:,
+          data_store:,
+          interval: 1,
+          task: Async::Task.current
+        )
+          task.async do
+            connect(collector_endpoint) do |client|
+              loop do
+                client.sync(data_store)
+                sleep(interval)
+              end
+            end
+          rescue Errno::EPIPE
+            Console.logger.error(self, "Broken pipe")
+          rescue Errno::ECONNREFUSED
+            Console.logger.error(self, "Connection refused")
+          end
+        end
+
+        sig { params(peer: Async::IO::Peer).void }
+        def initialize(peer)
+          wrapper = Wrapper.new
+          @packer = T.let(wrapper.packer(peer), MessagePack::Packer)
+        end
+
+        sig { params(data_store: DataStore, task: Async::Task).void }
+        def sync(data_store, task: Async::Task.current)
+          send(:store, data_store.store)
+        end
+
+        private
+
+        sig { params(args: T.untyped).void }
+        def send(*args)
+          @packer.write(args)
+          @packer.flush
+        end
+      end
+
+      class DataStore
+        class MetricStore
+          extend T::Sig
+
+          sig do
+            params(
+              store: ValueHash,
+              metric_name: Symbol,
+              metric_type: Symbol,
+              metric_settings: T::Hash[Symbol, T.untyped]
+            ).void
+          end
+          def initialize(store, metric_name:, metric_type:, metric_settings:)
+            @store = store
+            @metric_name = metric_name
+            @semaphore = T.let(Async::Semaphore.new, Async::Semaphore)
+          end
+
+          sig do
+            type_parameters(:T)
+              .params(block: T.proc.returns(T.type_parameter(:T)))
+              .returns(T.type_parameter(:T))
+          end
+          def synchronize(&block)
+            @semaphore.async { yield }.wait
+          end
+
+          sig do
+            params(
+              val: T.any(Integer, Float),
+              labels: T::Hash[Symbol, T.untyped]
+            ).void
+          end
+          def set(val:, labels: {})
+            @store.store(labels, val.to_f)
+          end
+
+          sig do
+            params(
+              by: T.any(Integer, Float),
+              labels: T::Hash[Symbol, T.untyped]
+            ).void
+          end
+          def increment(by: 1, labels: {})
+            @store.store(labels, @store.fetch(labels, 0.0) + by.to_f)
+          end
+
+          sig { params(labels: T::Hash[Symbol, T.untyped]).void }
+          def get(labels:)
+            @store.fetch(labels)
+          end
+        end
+
+        extend T::Sig
+
+        sig { returns(MetricHash) }
+        attr_reader :store
+
+        sig { void }
+        def initialize
+          @store = T.let(init_metric_hash, MetricHash)
+        end
+
+        sig do
+          params(
+            metric_name: Symbol,
+            metric_type: Symbol,
+            metric_settings: T::Hash[Symbol, T.untyped]
+          ).returns(MetricStore)
+        end
+        def for_metric(metric_name, metric_type:, metric_settings: {})
+          MetricStore.new(
+            T.must(@store[metric_name]),
+            metric_name:,
+            metric_type:,
+            metric_settings:
+          )
+        end
+
+        private
+
+        sig { returns(MetricHash) }
+        def init_metric_hash
+          Hash.new { |hash, metric_name| hash[metric_name] = init_value_hash }
+        end
+
+        sig { returns(ValueHash) }
+        def init_value_hash
+          Hash.new { |hash, labels| hash[labels] = 0 }
+        end
+      end
+    end
+  end
+end

data/lib/mayu/metrics.rb

@@ -0,0 +1,82 @@
+# typed: strict
+
+require "bundler/setup"
+require "async"
+require "async/container"
+require "async/semaphore"
+require "async/http"
+require "async/io/unix_endpoint"
+require "async/io/shared_endpoint"
+require "msgpack"
+require "nanoid"
+require "prometheus/client"
+require "prometheus/client/formats/text"
+
+require_relative "metrics/collector"
+require_relative "metrics/exporter"
+require_relative "metrics/reporter"
+
+module Mayu
+  module Metrics
+    InternalStore = T.type_alias { T::Hash[Symbol, MetricHash] }
+    MetricHash = T.type_alias { T::Hash[Symbol, ValueHash] }
+    ValueHash = T.type_alias { T::Hash[LabelsHash, Float] }
+    LabelsHash = T.type_alias { T::Hash[Symbol, T.untyped] }
+
+    class Wrapper < MessagePack::Factory
+      extend T::Sig
+
+      sig { void }
+      def initialize
+        super()
+
+        self.register_type(0x01, Symbol)
+      end
+    end
+
+    extend T::Sig
+
+    sig do
+      params(
+        container: Async::Container::Generic,
+        exporter_endpoint: Async::HTTP::Endpoint,
+        collector_endpoint: Async::IO::UNIXEndpoint,
+        block: T.proc.params(arg0: Prometheus::Client::Registry).void
+      ).void
+    end
+    def self.start_collect_and_export(
+      container,
+      exporter_endpoint:,
+      collector_endpoint:,
+      &block
+    )
+      collector = Metrics::Collector::Server.new(collector_endpoint)
+      collector.start
+
+      container.spawn(
+        name: "Metrics collector/exporter",
+        restart: true
+      ) do |instance|
+        Async do
+          internal_store = {}
+
+          Prometheus::Client.config.data_store =
+            Metrics::Collector::DataStore.new(internal_store)
+
+          registry = Prometheus::Client::Registry.new
+
+          yield registry
+
+          Metrics::Exporter::Server.setup(
+            endpoint: exporter_endpoint,
+            registry:
+          ).run
+
+          collector.run(internal_store)
+
+          instance.ready!
+        end
+      end
+    end
+  end
+end