progress dumping motd

got it working, need to clean up and test

libshpool/Cargo.toml

@@ -37,6 +37,8 @@ tracing = "0.1" # logging and performance monitoring facade
 bincode = "1" # serialization for the control protocol
 shpool_vt100 = "0.1.2" # terminal emulation for the scrollback buffer
 shell-words = "1" # parsing the -c/--cmd argument
+motd = "0.2.0" # getting the message-of-the-day
+terminfo = "0.8.0" # resolving terminal escape codes
 
 [dependencies.tracing-subscriber]
 version = "0.3"

libshpool/README.md

@@ -8,3 +8,14 @@ to an internal google version of the tool, but don't believe
 that telemetry belongs in an open-source tool. Other potential
 use-cases such as incorporating a shpool daemon into an
 IDE that hosts remote terminals could be imagined though.
+
+## Integrating
+
+In order to call libshpool, you must keep a few things in mind.
+In spirit, you just need to call `libshpool::run(libshpoo::Args::parse())`,
+but you need to take care of a few things manually.
+
+1. Handle the `version` subcommand. Since libshpool is a library, the output
+   will not be very good if the library handles the versioning.
+2. Depend on the `motd` crate and call `motd::handle_reexec()` in your `main`
+   function.

libshpool/src/config.rs

@@ -110,6 +110,9 @@ pub struct Config {
     /// verbatim except that the string '$SHPOOL_SESSION_NAME' will
     /// get replaced with the actual name of the shpool session.
     pub prompt_prefix: Option<String>,
+
+    /// Control when and how shpool will display the message of the day.
+    pub motd: Option<MotdDisplayMode>,
 }
 
 #[derive(Deserialize, Debug, Clone)]
@@ -140,6 +143,31 @@ pub enum SessionRestoreMode {
     Lines(u16),
 }
 
+#[derive(Deserialize, Debug, Clone, Default)]
+#[serde(rename_all = "lowercase")]
+pub enum MotdDisplayMode {
+    /// Never display the message of the day.
+    #[default]
+    Never,
+
+    /// Display the message of the day using the given program
+    /// as the pager. The pager will be invoked like `pager /tmp/motd.txt`,
+    /// and normal connection will only proceed once the pager has
+    /// exited.
+    ///
+    /// Display the message of the day each time a user attaches
+    /// (wether to a new session or reattaching to an existing session).
+    ///
+    /// `less` by default.
+    // Pager(String),
+
+    /// Just dump the message of the day directly to the screen.
+    /// Dumps are only performed when a new session is created.
+    /// There is no safe way to dump directly when reattaching,
+    /// so we don't attempt it.
+    Dump,
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

libshpool/src/daemon/control_codes.rs

@@ -0,0 +1,82 @@
+// Copyright 2024 Google LLC
+//
+// 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.
+
+//! The escape codes module provides an online (trie based) matcher
+//! to scan for escape codes we are interested in in the output of
+//! the subshell. For the moment, we just use this to scan for
+//! the ClearScreen code emitted by the prompt prefix injection shell
+//! code. We need to scan for this to avoid a race that can lead to
+//! the motd getting clobbered when in dump mode.
+
+use anyhow::{anyhow, Context};
+
+use super::trie::{TrieCursor, Trie};
+
+#[derive(Debug, Clone, Copy)]
+pub enum Code {
+    ClearScreen,
+}
+
+#[derive(Debug)]
+pub struct Matcher {
+    codes: Trie<u8, Code, Vec<Option<usize>>>,
+    codes_cursor: TrieCursor,
+}
+
+impl Matcher {
+    pub fn new(term_db: &terminfo::Database) -> anyhow::Result<Self> {
+        let clear_code = term_db.get::<terminfo::capability::ClearScreen>()
+            .ok_or(anyhow!("no clear screen code"))?;
+        let clear_code_bytes = clear_code.expand().to_vec().context("expanding clear code")?;
+
+        // TODO: delete
+        let clear_code_hex_bytes = clear_code_bytes.iter().map(|b| format!("{:02x}", b)).collect::<Vec<_>>().join(" ");
+        tracing::debug!("clear_code_hex_bytes={}", clear_code_hex_bytes);
+
+        let raw_bindings = vec![
+            // We need to scan for the clear code that gets emitted by the prompt prefix
+            // shell injection code so that we can make sure that the message of the day
+            // won't get clobbered immediately.
+            (clear_code_bytes, Code::ClearScreen),
+        ];
+        let mut codes = Trie::new();
+        for (raw_bytes, code) in raw_bindings.into_iter() {
+            codes.insert(raw_bytes.into_iter(), code);
+        }
+
+        Ok(Matcher {
+            codes,
+            codes_cursor: TrieCursor::Start,
+        })
+    }
+
+    pub fn transition(&mut self, byte: u8) -> Option<Code> {
+        let old_cursor = self.codes_cursor;
+        self.codes_cursor = self.codes.advance(self.codes_cursor, byte);
+        tracing::debug!("TRANSITION({:?}, {:02x}) -> {:?}", old_cursor, byte, self.codes_cursor); // TODO: delete
+        match self.codes_cursor {
+            TrieCursor::NoMatch => {
+                self.codes_cursor = TrieCursor::Start;
+                None
+            },
+            TrieCursor::Match { is_partial, .. } if !is_partial => {
+                let code = self.codes.get(self.codes_cursor).map(|c| *c);
+                tracing::debug!("TRANSITION MATCH: code={:?}", code);
+                self.codes_cursor = TrieCursor::Start;
+                code
+            }
+            _ => None,
+        }
+    }
+}

libshpool/src/daemon/keybindings.rs

@@ -48,11 +48,13 @@
 //! be singletons besides 'Ctrl' or of the form 'Ctrl-x' where
 //! x is some non-'Ctrl' key.
 
-use std::{collections::HashMap, fmt, hash};
+use std::{collections::HashMap, fmt};
 
 use anyhow::{anyhow, Context};
 use serde_derive::Deserialize;
 
+use super::trie::{Trie, TrieCursor, TrieTab};
+
 //
 // Keybindings table
 //
@@ -96,6 +98,20 @@ pub enum BindingResult {
 #[derive(Eq, PartialEq, Copy, Clone, Hash)]
 struct ChordAtom(u8);
 
+impl TrieTab<ChordAtom> for Vec<Option<usize>> {
+    fn new() -> Self {
+        vec![None; u8::MAX as usize]
+    }
+
+    fn get(&self, index: ChordAtom) -> Option<&usize> {
+        self[index.0 as usize].as_ref()
+    }
+
+    fn set(&mut self, index: ChordAtom, elem: usize) {
+        self[index.0 as usize] = Some(elem)
+    }
+}
+
 impl Bindings {
     /// new builds a bindings matching engine, parsing the given binding->action
     /// mapping and compiling it into the pair of tries that we use to perform
@@ -396,164 +412,6 @@ impl Lexer {
     }
 }
 
-//
-// Trie (used in both the parser and the execution engine)
-//
-
-#[derive(Debug)]
-struct Trie<Sym, V, TT> {
-    // The nodes which form the tree. The first node is the root
-    // node, afterwards the order is undefined.
-    nodes: Vec<TrieNode<Sym, V, TT>>,
-}
-
-#[derive(Eq, PartialEq, Copy, Clone, Debug)]
-enum TrieCursor {
-    /// A cursor to use to start a char-wise match
-    Start,
-    /// Represents a state in the middle or end of a match
-    Match { idx: usize, is_partial: bool },
-    /// A terminal state indicating a failure to match
-    NoMatch,
-}
-
-#[derive(Debug)]
-struct TrieNode<Sym, V, TT> {
-    // We need to store a phantom symbol here so we can have the
-    // Sym type parameter available for the TrieTab trait constraint
-    // in the impl block. Apologies for the type tetris.
-    phantom: std::marker::PhantomData<Sym>,
-    value: Option<V>,
-    tab: TT,
-}
-
-impl<Sym, V, TT> Trie<Sym, V, TT>
-where
-    TT: TrieTab<Sym>,
-    Sym: Copy,
-{
-    fn new() -> Self {
-        Trie { nodes: vec![TrieNode::new(None)] }
-    }
-
-    fn insert<Seq: Iterator<Item = Sym>>(&mut self, seq: Seq, value: V) {
-        let mut current_node = 0;
-        for sym in seq {
-            current_node = if let Some(next_node) = self.nodes[current_node].tab.get(sym) {
-                *next_node
-            } else {
-                let idx = self.nodes.len();
-                self.nodes.push(TrieNode::new(None));
-                self.nodes[current_node].tab.set(sym, idx);
-                idx
-            };
-        }
-        self.nodes[current_node].value = Some(value);
-    }
-
-    #[allow(dead_code)]
-    fn contains<Seq: Iterator<Item = Sym>>(&self, seq: Seq) -> bool {
-        let mut match_state = TrieCursor::Start;
-        for sym in seq {
-            match_state = self.advance(match_state, sym);
-            if let TrieCursor::NoMatch = match_state {
-                return false;
-            }
-        }
-        if let TrieCursor::Start = match_state {
-            return self.nodes[0].value.is_some();
-        }
-
-        if let TrieCursor::Match { is_partial, .. } = match_state { !is_partial } else { false }
-    }
-
-    fn advance(&self, cursor: TrieCursor, sym: Sym) -> TrieCursor {
-        let node = match cursor {
-            TrieCursor::Start => &self.nodes[0],
-            TrieCursor::Match { idx, .. } => &self.nodes[idx],
-            TrieCursor::NoMatch => return TrieCursor::NoMatch,
-        };
-
-        if let Some(idx) = node.tab.get(sym) {
-            TrieCursor::Match { idx: *idx, is_partial: self.nodes[*idx].value.is_none() }
-        } else {
-            TrieCursor::NoMatch
-        }
-    }
-
-    fn get(&self, cursor: TrieCursor) -> Option<&V> {
-        if let TrieCursor::Match { idx, .. } = cursor {
-            self.nodes[idx].value.as_ref()
-        } else {
-            None
-        }
-    }
-}
-
-impl<Sym, V, TT> TrieNode<Sym, V, TT>
-where
-    TT: TrieTab<Sym>,
-{
-    fn new(value: Option<V>) -> Self {
-        TrieNode { phantom: std::marker::PhantomData, value, tab: TT::new() }
-    }
-}
-
-/// The backing table the trie uses to associate symbols with state
-/// indexes. This is basically std::ops::IndexMut plus a new function.
-/// We can't just make this a sub-trait of IndexMut because u8 does
-/// not implement IndexMut for vectors.
-trait TrieTab<Idx> {
-    fn new() -> Self;
-    fn get(&self, index: Idx) -> Option<&usize>;
-    fn set(&mut self, index: Idx, elem: usize);
-}
-
-impl<Sym> TrieTab<Sym> for HashMap<Sym, usize>
-where
-    Sym: hash::Hash + Eq + PartialEq,
-{
-    fn new() -> Self {
-        HashMap::new()
-    }
-
-    fn get(&self, index: Sym) -> Option<&usize> {
-        self.get(&index)
-    }
-
-    fn set(&mut self, index: Sym, elem: usize) {
-        self.insert(index, elem);
-    }
-}
-
-impl TrieTab<u8> for Vec<Option<usize>> {
-    fn new() -> Self {
-        vec![None; u8::MAX as usize]
-    }
-
-    fn get(&self, index: u8) -> Option<&usize> {
-        self[index as usize].as_ref()
-    }
-
-    fn set(&mut self, index: u8, elem: usize) {
-        self[index as usize] = Some(elem)
-    }
-}
-
-impl TrieTab<ChordAtom> for Vec<Option<usize>> {
-    fn new() -> Self {
-        vec![None; u8::MAX as usize]
-    }
-
-    fn get(&self, index: ChordAtom) -> Option<&usize> {
-        self[index.0 as usize].as_ref()
-    }
-
-    fn set(&mut self, index: ChordAtom, elem: usize) {
-        self[index.0 as usize] = Some(elem)
-    }
-}
-
 //
 // Data Tables
 //

libshpool/src/daemon/mod.rs

@@ -28,6 +28,9 @@ mod shell;
 mod signals;
 mod systemd;
 mod ttl_reaper;
+mod show_motd;
+mod trie;
+mod control_codes;
 
 #[instrument(skip_all)]
 pub fn run(
@@ -39,7 +42,7 @@ pub fn run(
     info!("\n\n======================== STARTING DAEMON ============================\n\n");
 
     let config = config::read_config(&config_file)?;
-    let server = server::Server::new(config, hooks, runtime_dir);
+    let server = server::Server::new(config, hooks, runtime_dir)?;
 
     let (cleanup_socket, listener) = match systemd::activation_socket() {
         Ok(l) => {