{-# LANGUAGE DuplicateRecordFields #-}

module Cachix.API.WebSocketSubprotocol where

import qualified Control.Concurrent.Async as Async
import qualified Control.Concurrent.STM.TMQueue as TMQueue
import qualified Control.Exception.Safe as Safe
import qualified Data.Aeson as Aeson
import Data.Time (UTCTime)
import Data.UUID (UUID)
import qualified Network.WebSockets as WS
import Protolude

data Message cmd = Message
  { forall cmd. Message cmd -> Text
method :: Text,
    forall cmd. Message cmd -> cmd
command :: cmd,
    forall cmd. Message cmd -> Maybe UUID
agent :: Maybe UUID,
    forall cmd. Message cmd -> UUID
id :: UUID
  }
  deriving (Int -> Message cmd -> ShowS
forall cmd. Show cmd => Int -> Message cmd -> ShowS
forall cmd. Show cmd => [Message cmd] -> ShowS
forall cmd. Show cmd => Message cmd -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Message cmd] -> ShowS
$cshowList :: forall cmd. Show cmd => [Message cmd] -> ShowS
show :: Message cmd -> String
$cshow :: forall cmd. Show cmd => Message cmd -> String
showsPrec :: Int -> Message cmd -> ShowS
$cshowsPrec :: forall cmd. Show cmd => Int -> Message cmd -> ShowS
Show, Message cmd -> Message cmd -> Bool
forall cmd. Eq cmd => Message cmd -> Message cmd -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Message cmd -> Message cmd -> Bool
$c/= :: forall cmd. Eq cmd => Message cmd -> Message cmd -> Bool
== :: Message cmd -> Message cmd -> Bool
$c== :: forall cmd. Eq cmd => Message cmd -> Message cmd -> Bool
Eq, forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall cmd x. Rep (Message cmd) x -> Message cmd
forall cmd x. Message cmd -> Rep (Message cmd) x
$cto :: forall cmd x. Rep (Message cmd) x -> Message cmd
$cfrom :: forall cmd x. Message cmd -> Rep (Message cmd) x
Generic, forall cmd. FromJSON cmd => Value -> Parser [Message cmd]
forall cmd. FromJSON cmd => Value -> Parser (Message cmd)
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [Message cmd]
$cparseJSONList :: forall cmd. FromJSON cmd => Value -> Parser [Message cmd]
parseJSON :: Value -> Parser (Message cmd)
$cparseJSON :: forall cmd. FromJSON cmd => Value -> Parser (Message cmd)
Aeson.FromJSON, forall cmd. ToJSON cmd => [Message cmd] -> Encoding
forall cmd. ToJSON cmd => [Message cmd] -> Value
forall cmd. ToJSON cmd => Message cmd -> Encoding
forall cmd. ToJSON cmd => Message cmd -> Value
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
toEncodingList :: [Message cmd] -> Encoding
$ctoEncodingList :: forall cmd. ToJSON cmd => [Message cmd] -> Encoding
toJSONList :: [Message cmd] -> Value
$ctoJSONList :: forall cmd. ToJSON cmd => [Message cmd] -> Value
toEncoding :: Message cmd -> Encoding
$ctoEncoding :: forall cmd. ToJSON cmd => Message cmd -> Encoding
toJSON :: Message cmd -> Value
$ctoJSON :: forall cmd. ToJSON cmd => Message cmd -> Value
Aeson.ToJSON)

data Cache = Cache
  { Cache -> Text
cacheName :: Text,
    Cache -> Text
publicKey :: Text,
    Cache -> Bool
isPublic :: Bool
  }
  deriving (Int -> Cache -> ShowS
[Cache] -> ShowS
Cache -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Cache] -> ShowS
$cshowList :: [Cache] -> ShowS
show :: Cache -> String
$cshow :: Cache -> String
showsPrec :: Int -> Cache -> ShowS
$cshowsPrec :: Int -> Cache -> ShowS
Show, Cache -> Cache -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Cache -> Cache -> Bool
$c/= :: Cache -> Cache -> Bool
== :: Cache -> Cache -> Bool
$c== :: Cache -> Cache -> Bool
Eq, forall x. Rep Cache x -> Cache
forall x. Cache -> Rep Cache x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Cache x -> Cache
$cfrom :: forall x. Cache -> Rep Cache x
Generic, Value -> Parser [Cache]
Value -> Parser Cache
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [Cache]
$cparseJSONList :: Value -> Parser [Cache]
parseJSON :: Value -> Parser Cache
$cparseJSON :: Value -> Parser Cache
Aeson.FromJSON, [Cache] -> Encoding
[Cache] -> Value
Cache -> Encoding
Cache -> Value
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
toEncodingList :: [Cache] -> Encoding
$ctoEncodingList :: [Cache] -> Encoding
toJSONList :: [Cache] -> Value
$ctoJSONList :: [Cache] -> Value
toEncoding :: Cache -> Encoding
$ctoEncoding :: Cache -> Encoding
toJSON :: Cache -> Value
$ctoJSON :: Cache -> Value
Aeson.ToJSON)

data AgentInformation = AgentInformation
  { AgentInformation -> Maybe Cache
cache :: Maybe Cache,
    AgentInformation -> UUID
id :: UUID
  }
  deriving (Int -> AgentInformation -> ShowS
[AgentInformation] -> ShowS
AgentInformation -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AgentInformation] -> ShowS
$cshowList :: [AgentInformation] -> ShowS
show :: AgentInformation -> String
$cshow :: AgentInformation -> String
showsPrec :: Int -> AgentInformation -> ShowS
$cshowsPrec :: Int -> AgentInformation -> ShowS
Show, AgentInformation -> AgentInformation -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AgentInformation -> AgentInformation -> Bool
$c/= :: AgentInformation -> AgentInformation -> Bool
== :: AgentInformation -> AgentInformation -> Bool
$c== :: AgentInformation -> AgentInformation -> Bool
Eq, forall x. Rep AgentInformation x -> AgentInformation
forall x. AgentInformation -> Rep AgentInformation x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep AgentInformation x -> AgentInformation
$cfrom :: forall x. AgentInformation -> Rep AgentInformation x
Generic, Value -> Parser [AgentInformation]
Value -> Parser AgentInformation
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [AgentInformation]
$cparseJSONList :: Value -> Parser [AgentInformation]
parseJSON :: Value -> Parser AgentInformation
$cparseJSON :: Value -> Parser AgentInformation
Aeson.FromJSON, [AgentInformation] -> Encoding
[AgentInformation] -> Value
AgentInformation -> Encoding
AgentInformation -> Value
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
toEncodingList :: [AgentInformation] -> Encoding
$ctoEncodingList :: [AgentInformation] -> Encoding
toJSONList :: [AgentInformation] -> Value
$ctoJSONList :: [AgentInformation] -> Value
toEncoding :: AgentInformation -> Encoding
$ctoEncoding :: AgentInformation -> Encoding
toJSON :: AgentInformation -> Value
$ctoJSON :: AgentInformation -> Value
Aeson.ToJSON)

data DeploymentDetails = DeploymentDetails
  { DeploymentDetails -> Text
storePath :: Text,
    DeploymentDetails -> UUID
id :: UUID,
    DeploymentDetails -> Int64
index :: Int64,
    DeploymentDetails -> Maybe Text
rollbackScript :: Maybe Text
  }
  deriving (Int -> DeploymentDetails -> ShowS
[DeploymentDetails] -> ShowS
DeploymentDetails -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DeploymentDetails] -> ShowS
$cshowList :: [DeploymentDetails] -> ShowS
show :: DeploymentDetails -> String
$cshow :: DeploymentDetails -> String
showsPrec :: Int -> DeploymentDetails -> ShowS
$cshowsPrec :: Int -> DeploymentDetails -> ShowS
Show, DeploymentDetails -> DeploymentDetails -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DeploymentDetails -> DeploymentDetails -> Bool
$c/= :: DeploymentDetails -> DeploymentDetails -> Bool
== :: DeploymentDetails -> DeploymentDetails -> Bool
$c== :: DeploymentDetails -> DeploymentDetails -> Bool
Eq, forall x. Rep DeploymentDetails x -> DeploymentDetails
forall x. DeploymentDetails -> Rep DeploymentDetails x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep DeploymentDetails x -> DeploymentDetails
$cfrom :: forall x. DeploymentDetails -> Rep DeploymentDetails x
Generic, Value -> Parser [DeploymentDetails]
Value -> Parser DeploymentDetails
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [DeploymentDetails]
$cparseJSONList :: Value -> Parser [DeploymentDetails]
parseJSON :: Value -> Parser DeploymentDetails
$cparseJSON :: Value -> Parser DeploymentDetails
Aeson.FromJSON, [DeploymentDetails] -> Encoding
[DeploymentDetails] -> Value
DeploymentDetails -> Encoding
DeploymentDetails -> Value
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
toEncodingList :: [DeploymentDetails] -> Encoding
$ctoEncodingList :: [DeploymentDetails] -> Encoding
toJSONList :: [DeploymentDetails] -> Value
$ctoJSONList :: [DeploymentDetails] -> Value
toEncoding :: DeploymentDetails -> Encoding
$ctoEncoding :: DeploymentDetails -> Encoding
toJSON :: DeploymentDetails -> Value
$ctoJSON :: DeploymentDetails -> Value
Aeson.ToJSON)

data BackendCommand
  = Deployment DeploymentDetails
  | AgentRegistered AgentInformation
  deriving (Int -> BackendCommand -> ShowS
[BackendCommand] -> ShowS
BackendCommand -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BackendCommand] -> ShowS
$cshowList :: [BackendCommand] -> ShowS
show :: BackendCommand -> String
$cshow :: BackendCommand -> String
showsPrec :: Int -> BackendCommand -> ShowS
$cshowsPrec :: Int -> BackendCommand -> ShowS
Show, BackendCommand -> BackendCommand -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: BackendCommand -> BackendCommand -> Bool
$c/= :: BackendCommand -> BackendCommand -> Bool
== :: BackendCommand -> BackendCommand -> Bool
$c== :: BackendCommand -> BackendCommand -> Bool
Eq, forall x. Rep BackendCommand x -> BackendCommand
forall x. BackendCommand -> Rep BackendCommand x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep BackendCommand x -> BackendCommand
$cfrom :: forall x. BackendCommand -> Rep BackendCommand x
Generic, Value -> Parser [BackendCommand]
Value -> Parser BackendCommand
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [BackendCommand]
$cparseJSONList :: Value -> Parser [BackendCommand]
parseJSON :: Value -> Parser BackendCommand
$cparseJSON :: Value -> Parser BackendCommand
Aeson.FromJSON, [BackendCommand] -> Encoding
[BackendCommand] -> Value
BackendCommand -> Encoding
BackendCommand -> Value
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
toEncodingList :: [BackendCommand] -> Encoding
$ctoEncodingList :: [BackendCommand] -> Encoding
toJSONList :: [BackendCommand] -> Value
$ctoJSONList :: [BackendCommand] -> Value
toEncoding :: BackendCommand -> Encoding
$ctoEncoding :: BackendCommand -> Encoding
toJSON :: BackendCommand -> Value
$ctoJSON :: BackendCommand -> Value
Aeson.ToJSON)

data AgentCommand
  = DeploymentStarted {AgentCommand -> UUID
id :: UUID, AgentCommand -> UTCTime
time :: UTCTime, AgentCommand -> Maybe Int64
closureSize :: Maybe Int64}
  | DeploymentFinished {id :: UUID, time :: UTCTime, AgentCommand -> Bool
hasSucceeded :: Bool}
  deriving (Int -> AgentCommand -> ShowS
[AgentCommand] -> ShowS
AgentCommand -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AgentCommand] -> ShowS
$cshowList :: [AgentCommand] -> ShowS
show :: AgentCommand -> String
$cshow :: AgentCommand -> String
showsPrec :: Int -> AgentCommand -> ShowS
$cshowsPrec :: Int -> AgentCommand -> ShowS
Show, AgentCommand -> AgentCommand -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AgentCommand -> AgentCommand -> Bool
$c/= :: AgentCommand -> AgentCommand -> Bool
== :: AgentCommand -> AgentCommand -> Bool
$c== :: AgentCommand -> AgentCommand -> Bool
Eq, forall x. Rep AgentCommand x -> AgentCommand
forall x. AgentCommand -> Rep AgentCommand x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep AgentCommand x -> AgentCommand
$cfrom :: forall x. AgentCommand -> Rep AgentCommand x
Generic, Value -> Parser [AgentCommand]
Value -> Parser AgentCommand
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [AgentCommand]
$cparseJSONList :: Value -> Parser [AgentCommand]
parseJSON :: Value -> Parser AgentCommand
$cparseJSON :: Value -> Parser AgentCommand
Aeson.FromJSON, [AgentCommand] -> Encoding
[AgentCommand] -> Value
AgentCommand -> Encoding
AgentCommand -> Value
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
toEncodingList :: [AgentCommand] -> Encoding
$ctoEncodingList :: [AgentCommand] -> Encoding
toJSONList :: [AgentCommand] -> Value
$ctoJSONList :: [AgentCommand] -> Value
toEncoding :: AgentCommand -> Encoding
$ctoEncoding :: AgentCommand -> Encoding
toJSON :: AgentCommand -> Value
$ctoJSON :: AgentCommand -> Value
Aeson.ToJSON)

parseMessage :: Aeson.FromJSON cmd => ByteString -> Either Text (Message cmd)
parseMessage :: forall cmd. FromJSON cmd => ByteString -> Either Text (Message cmd)
parseMessage = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a b. ConvertText a b => a -> b
toS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. FromJSON a => ByteString -> Either String a
Aeson.eitherDecodeStrict'

sendMessage :: Aeson.ToJSON cmd => WS.Connection -> Message cmd -> IO ()
sendMessage :: forall cmd. ToJSON cmd => Connection -> Message cmd -> IO ()
sendMessage Connection
connection Message cmd
cmd =
  forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendTextData Connection
connection (forall a. ToJSON a => a -> ByteString
Aeson.encode Message cmd
cmd)

-- | Receive and process messages in parallel.
--
-- Note: This will not rethrow the 'CloseRequest' exception!
--
-- TODO: use Async.replicateConcurrently
receiveDataConcurrently :: WS.Connection -> (ByteString -> IO ()) -> IO ()
receiveDataConcurrently :: Connection -> (ByteString -> IO ()) -> IO ()
receiveDataConcurrently Connection
connection ByteString -> IO ()
action =
  do
    TMQueue ByteString
queue <- forall a. STM a -> IO a
atomically forall a. STM (TMQueue a)
TMQueue.newTMQueue
    forall a b. IO a -> (Async a -> IO b) -> IO b
Async.withAsync (TMQueue ByteString -> IO ()
consumer TMQueue ByteString
queue) (forall {a}. WebSocketsData a => TMQueue a -> Async () -> IO ()
producer TMQueue ByteString
queue)
  where
    producer :: TMQueue a -> Async () -> IO ()
producer TMQueue a
queue Async ()
consumerThread =
      forall {b}. IO b
loop
        forall a b. IO a -> IO b -> IO a
`finally` forall a. TMQueue a -> Async () -> IO ()
closeGracefully TMQueue a
queue Async ()
consumerThread
        forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> (e -> m a) -> m a
`Safe.catch` ConnectionException -> IO ()
closeRequest
      where
        loop :: IO b
loop = do
          a
payload <- forall a. WebSocketsData a => Connection -> IO a
WS.receiveData Connection
connection
          forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMQueue a -> a -> STM ()
TMQueue.writeTMQueue TMQueue a
queue a
payload
          IO b
loop

    consumer :: TMQueue ByteString -> IO ()
consumer TMQueue ByteString
queue = IO ()
loop
      where
        loop :: IO ()
loop = do
          Maybe ByteString
payload <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMQueue a -> STM (Maybe a)
TMQueue.readTMQueue TMQueue ByteString
queue
          case Maybe ByteString
payload of
            Maybe ByteString
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Just ByteString
msg -> ByteString -> IO ()
action ByteString
msg forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> IO ()
loop

    -- Close the queue and wait for the consumer finish processing messages.
    closeGracefully :: TMQueue.TMQueue a -> Async.Async () -> IO ()
    closeGracefully :: forall a. TMQueue a -> Async () -> IO ()
closeGracefully TMQueue a
queue Async ()
consumerThread = do
      forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMQueue a -> STM ()
TMQueue.closeTMQueue TMQueue a
queue
      forall a. Async a -> IO a
Async.wait Async ()
consumerThread

    closeRequest :: WS.ConnectionException -> IO ()
    closeRequest :: ConnectionException -> IO ()
closeRequest (WS.CloseRequest Word16
_ ByteString
_) = forall (m :: * -> *) a. Monad m => a -> m a
return ()
    closeRequest ConnectionException
e = forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO ConnectionException
e

data Log = Log
  { Log -> Text
line :: Text,
    Log -> UTCTime
time :: UTCTime
  }
  deriving (forall x. Rep Log x -> Log
forall x. Log -> Rep Log x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Log x -> Log
$cfrom :: forall x. Log -> Rep Log x
Generic, Int -> Log -> ShowS
[Log] -> ShowS
Log -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Log] -> ShowS
$cshowList :: [Log] -> ShowS
show :: Log -> String
$cshow :: Log -> String
showsPrec :: Int -> Log -> ShowS
$cshowsPrec :: Int -> Log -> ShowS
Show, [Log] -> Encoding
[Log] -> Value
Log -> Encoding
Log -> Value
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
toEncodingList :: [Log] -> Encoding
$ctoEncodingList :: [Log] -> Encoding
toJSONList :: [Log] -> Value
$ctoJSONList :: [Log] -> Value
toEncoding :: Log -> Encoding
$ctoEncoding :: Log -> Encoding
toJSON :: Log -> Value
$ctoJSON :: Log -> Value
Aeson.ToJSON, Value -> Parser [Log]
Value -> Parser Log
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [Log]
$cparseJSONList :: Value -> Parser [Log]
parseJSON :: Value -> Parser Log
$cparseJSON :: Value -> Parser Log
Aeson.FromJSON)