websockets and better (empty) response messages

(tough mostly untested)

2 files changed, 78 insertions, 57 deletions

diff --git a/lib/API.hs b/lib/API.hs
index 845ad06..8d352e1 100644
--- a/lib/API.hs
+++ b/lib/API.hs
@@ -22,9 +22,11 @@ import           Servant                   (Application,
                                             Server, err401, err404, serve,
                                             throwError, type (:>))
 import           Servant.API               (Capture, FromHttpApiData, Get, JSON,
-                                            Post, QueryParam, ReqBody,
-                                            type (:<|>) ((:<|>)))
+                                            NoContent, Post, QueryParam,
+                                            ReqBody, type (:<|>) ((:<|>)))
+import           Servant.API.WebSocket     (WebSocket)
 import           Servant.GTFS.Realtime     (Proto)
+import           Servant.Swagger           (HasSwagger (..))
 
 import           Data.UUID                 (UUID)
 
@@ -36,7 +38,8 @@ type API = "stations" :> Get '[JSON] (Map StationID Station)
   -- TODO: perhaps require a first ping for registration?
   :<|> "train" :> "register" :> Capture "Trip ID" TripID :> Post '[JSON] Token
   -- TODO: perhaps a websocket instead?
-  :<|> "train" :> "ping" :> ReqBody '[JSON] TripPing :> Post '[JSON] ()
+  :<|> "train" :> "ping" :> ReqBody '[JSON] TripPing :> Post '[JSON] NoContent
+  :<|> "train" :> "ping" :> "ws" :> WebSocket
   -- debug things
   :<|> "debug" :> "state" :> Get '[JSON] (Map Token [TripPing])
   :<|> "gtfs" :> GtfsRealtimeAPI
@@ -51,12 +54,12 @@ type GtfsRealtimeAPI = "servicealerts" :> Get '[Proto] FeedMessage
 -- | Admin API used for short-term timetable changes etc. ("leitstelle")
 type AdminAPI =
   "trip" :> "announce" :> ReqBody '[JSON] Announcement :> Post '[JSON] UUID
-  :<|> "trip" :> "announce" :> "delete" :> Capture "Announcement ID" UUID :> Post '[JSON] ()
-  :<|> "trip" :> "date" :> "add" :> Capture "Trip ID" TripID :> Capture "day" Day :> Post '[JSON] ()
-  :<|> "trip" :> "date" :> "cancel" :> Capture "Trip ID" TripID :> Capture "day" Day :> Post '[JSON] ()
+  :<|> "trip" :> "announce" :> "delete" :> Capture "Announcement ID" UUID :> Post '[JSON] NoContent
+  :<|> "trip" :> "date" :> "add" :> Capture "Trip ID" TripID :> Capture "day" Day :> Post '[JSON] NoContent
+  :<|> "trip" :> "date" :> "cancel" :> Capture "Trip ID" TripID :> Capture "day" Day :> Post '[JSON] NoContent
 -- TODO for this to be useful there ought to be a half-deep Trip type
 -- (that has stops but not shapes)
-  :<|> "extraordinary" :> "trip" :> ReqBody '[JSON] (Trip Deep Shallow) :> Post '[JSON] ()
+  :<|> "extraordinary" :> "trip" :> ReqBody '[JSON] (Trip Deep Shallow) :> Post '[JSON] NoContent
 
 
 -- | The server's API with an additional debug route for accessing the specification
@@ -70,8 +73,9 @@ type CompleteAPI = "debug" :> "openapi" :> Get '[JSON] Swagger
 instance ToParamSchema (Maybe UTCTime) where
   toParamSchema _ = toParamSchema (Proxy @UTCTime)
 
-
-
+-- TODO write something useful here! (and if it's just "hey this is some websocket thingie")
+instance HasSwagger WebSocket where
+  toSwagger _ = toSwagger (Proxy @(Post '[JSON] NoContent))
 {-
 TODO:
 there should be a basic API allowing the questions:
diff --git a/lib/Server.hs b/lib/Server.hs
index 7802911..8a6022c 100644
--- a/lib/Server.hs
+++ b/lib/Server.hs
@@ -10,58 +10,61 @@
 
 -- Implementation of the API. This module is the main point of the program.
 module Server (application) where
-import           Conduit                        (MonadTrans (lift), ResourceT)
+import           Conduit                     (MonadTrans (lift), ResourceT)
 import           Control.Concurrent.STM
-import           Control.Monad                  (when)
-import           Control.Monad.Extra            (maybeM, whenM)
-import           Control.Monad.IO.Class         (MonadIO (liftIO))
-import           Control.Monad.Logger.CallStack (NoLoggingT)
-import           Control.Monad.Reader           (forM)
-import           Control.Monad.Trans.Maybe      (MaybeT (..))
-import           Data.Aeson                     (FromJSON (parseJSON),
-                                                 ToJSON (toJSON), ToJSONKey,
-                                                 genericParseJSON,
-                                                 genericToJSON)
-import qualified Data.Aeson                     as A
-import           Data.Coerce                    (coerce)
-import           Data.Functor                   ((<&>))
-import           Data.Map                       (Map)
-import qualified Data.Map                       as M
-import           Data.Pool                      (Pool)
-import           Data.Proxy                     (Proxy (Proxy))
-import           Data.Swagger                   hiding (get, delete)
-import           Data.Text                      (Text)
-import           Data.Time                      (NominalDiffTime,
-                                                 UTCTime (utctDay), addUTCTime,
-                                                 dayOfWeek, diffUTCTime,
-                                                 getCurrentTime, nominalDay)
-import           Data.UUID                      (UUID)
-import qualified Data.UUID                      as UUID
-import qualified Data.UUID.V4                   as UUID
-import           Data.Vector                    (Vector)
-import qualified Data.Vector                    as V
+import           Control.Monad               (forever, void, when)
+import           Control.Monad.Extra         (maybeM, whenM)
+import           Control.Monad.IO.Class      (MonadIO (liftIO))
+import           Control.Monad.Logger        (NoLoggingT, logWarnN,
+                                              runStderrLoggingT)
+import           Control.Monad.Reader        (forM)
+import           Control.Monad.Trans.Maybe   (MaybeT (..))
+import           Data.Aeson                  (FromJSON (parseJSON),
+                                              ToJSON (toJSON), ToJSONKey,
+                                              genericParseJSON, genericToJSON)
+import qualified Data.Aeson                  as A
+import qualified Data.ByteString.Char8       as C8
+import           Data.Coerce                 (coerce)
+import           Data.Functor                ((<&>))
+import           Data.Map                    (Map)
+import qualified Data.Map                    as M
+import           Data.Pool                   (Pool)
+import           Data.Proxy                  (Proxy (Proxy))
+import           Data.Swagger                hiding (delete, get)
+import           Data.Text                   (Text)
+import           Data.Time                   (NominalDiffTime,
+                                              UTCTime (utctDay), addUTCTime,
+                                              dayOfWeek, diffUTCTime,
+                                              getCurrentTime, nominalDay)
+import           Data.UUID                   (UUID)
+import qualified Data.UUID                   as UUID
+import qualified Data.UUID.V4                as UUID
+import           Data.Vector                 (Vector)
+import qualified Data.Vector                 as V
 import           Database.Persist
 import           Database.Persist.Postgresql
-import           GHC.Generics                   (Generic)
+import           Fmt                         ((+|), (|+))
+import           GHC.Generics                (Generic)
 import           GTFS
-import           Servant                        (Application,
-                                                 FromHttpApiData (parseUrlPiece),
-                                                 Server, err401, err404, serve,
-                                                 throwError)
-import           Servant.API                    (Capture, FromHttpApiData, Get,
-                                                 JSON, Post, ReqBody,
-                                                 type (:<|>) ((:<|>)))
-import           Servant.Docs                   (DocCapture (..),
-                                                 DocQueryParam (..),
-                                                 ParamKind (..), ToCapture (..),
-                                                 ToParam (..))
-import           Servant.Server                 (Handler)
-import           Servant.Swagger                (toSwagger)
-import           Web.PathPieces                 (PathPiece)
+import qualified Network.WebSockets          as WS
+import           Servant                     (Application,
+                                              FromHttpApiData (parseUrlPiece),
+                                              Server, err401, err404, serve,
+                                              throwError)
+import           Servant.API                 (Capture, FromHttpApiData, Get,
+                                              JSON, NoContent (..), Post,
+                                              ReqBody, type (:<|>) ((:<|>)))
+import           Servant.Docs                (DocCapture (..),
+                                              DocQueryParam (..),
+                                              ParamKind (..), ToCapture (..),
+                                              ToParam (..))
+import           Servant.Server              (Handler)
+import           Servant.Swagger             (toSwagger)
+import           Web.PathPieces              (PathPiece)
 
 import           API
 import           Persist
-import           Server.GTFS_RT                  (gtfsRealtimeServer)
+import           Server.GTFS_RT              (gtfsRealtimeServer)
 
 application :: GTFS -> Pool SqlBackend -> IO Application
 application gtfs dbpool = do
@@ -77,7 +80,7 @@ doMigration pool = runSql pool $
 
 server :: GTFS -> Pool SqlBackend -> Server CompleteAPI
 server gtfs@GTFS{..} dbpool = handleDebugAPI :<|> handleStations :<|> handleTimetable :<|> handleTrip
-  :<|> handleRegister :<|> handleTripPing :<|> handleDebugState :<|> gtfsRealtimeServer gtfs dbpool
+  :<|> handleRegister :<|> handleTripPing :<|> handleWS :<|> handleDebugState :<|> gtfsRealtimeServer gtfs dbpool
   :<|> adminServer gtfs dbpool
   where handleStations = pure stations
         handleTimetable station maybeDay = do
@@ -99,7 +102,19 @@ server gtfs@GTFS{..} dbpool = handleDebugAPI :<|> handleStations :<|> handleTime
           checkTokenValid dbpool (coerce $ tripPingToken ping)
           -- TODO: are these always inserted in order?
           runSql dbpool $ insert ping
-          pure ()
+          pure NoContent
+        handleWS conn = do
+          -- TODO test this!!
+          liftIO $ WS.forkPingThread conn 30
+          forever $ do
+            msg <- liftIO $ WS.receiveData conn
+            case A.eitherDecode msg of
+              Left err -> runStderrLoggingT $ do
+                logWarnN ("stray websocket message: "+|show msg|+" (could not decode: "+|err|+")")
+                liftIO $ WS.sendClose conn (C8.pack err)
+              Right ping -> do
+                checkTokenValid dbpool (coerce $ tripPingToken ping)
+                void $ runSql dbpool $ insert ping
         handleDebugState = do
           now <- liftIO getCurrentTime
           runSql dbpool $ do
@@ -120,13 +135,15 @@ adminServer gtfs dbpool =
           pure uuid
         delAnnounce uuid = runSql dbpool $ do
           delete (AnnouncementKey uuid)
+          pure NoContent
         modTripDate one two tripId day = runSql dbpool $ do
           getBy (TripAndDay tripId day) >>= \case
             Just (Entity key (ScheduleAmendment _ _ status)) -> do
               when (status == two) $ delete key
+              pure NoContent
             Nothing -> do
               insert (ScheduleAmendment tripId day one)
-              pure ()
+              pure NoContent
         extraTrip = error "unimplemented!"