State machine for handling Telnet sequences for an FTP server

Question

I've asked Code Climate to generate metrics for the ftpd Ruby gem. It correctly identified the God class; I know what to do about that. But one of the smaller classes has me stumped. This is telnet.rb:

# -*- ruby encoding: us-ascii -*-

module Ftpd

  # Handle the limited processing of Telnet sequences required by the
  # FTP RFCs.
  #
  # Telnet option processing is quite complex, but we need do only a
  # simple subset of it, since we can disagree with any request by the
  # client to turn on an option (RFC-1123 4.1.2.12).  Adhering to
  # RFC-1143 ("The Q Method of Implementing TELNET Option Negiation"),
  # and supporting only what's needed to keep all options turned off:
  #
  # * Reply to WILL sequence with DONT sequence
  # * Reply to DO sequence with WONT sequence
  # * Ignore WONT sequence
  # * Ignore DONT sequence
  #
  # We also handle the "interrupt process" and "data mark" sequences,
  # which the client sends before the ABORT command, by ignoring them.
  #
  # All Telnet sequence start with an IAC, followed by at least one
  # character.  Here are the sequences we care about:
  #
  #     SEQUENCE             CODES
  #     -----------------    --------------------
  #     WILL                 IAC WILL option-code
  #     WONT                 IAC WONT option-code
  #     DO                   IAC DO option-code
  #     DONT                 IAC DONT option-code
  #     escaped 255          IAC IAC
  #     interrupt process    IAC IP
  #     data mark            IAC DM
  #
  # Any pathalogical sequence (e.g. IAC + \x01), or any sequence we
  # don't recognize, we pass through.

  class Telnet

    # The command with recognized Telnet sequences removed

    attr_reader :plain

    # Any Telnet sequences to send

    attr_reader :reply

    # Create a new instance with a command that may contain Telnet
    # sequences.
    # @param command [String]

    def initialize(command)
      telnet_state_machine command
    end

    private

    module Codes
      IAC  = 255.chr    # 0xff
      DONT = 254.chr    # 0xfe
      DO   = 253.chr    # 0xfd
      WONT = 252.chr    # 0xfc
      WILL = 251.chr    # 0xfb
      IP   = 244.chr    # 0xf4
      DM   = 242.chr    # 0xf2
    end
    include Codes

    def telnet_state_machine (command)
      @plain = ''
      @reply = ''
      state = :idle
      command.each_char do |c|
        case state
        when :idle
          if c == IAC
            state = :iac
          else
            @plain << c
          end
        when :iac
          case c
          when IAC
            @plain << c
            state = :idle
          when WILL
            state = :will
          when WONT
            state = :wont
          when DO
            state = :do
          when DONT
            state = :dont
          when IP
            state = :idle
          when DM
            state = :idle
          else
            @plain << IAC + c
            state = :idle
          end
        when :will
          @reply << IAC + DONT + c
          state = :idle
        when :wont
          state = :idle
        when :do
          @reply << IAC + WONT + c
          state = :idle
        when :dont
          state = :idle
        else
          raise "Unknown state #{state.inspect}"
        end
      end
    end

  end
end

Code Climate does not like the complexity of #telnet_state_machine. I agree, but I don't know how to reduce the complexity without also making the state machine harder to follow. State machines never seem to be all that readable as it is. What would you suggest?

Note: If you want to try some refactoring, this class has rspec test coverage. Just do "git clone" of the ftpd project.

Answer 1

EDIT

the more I think about it, the more it seems obvious that a parser would be more adapted to this problem than a state machine. Unfortunately, I don't know much about writing parsers...

ORIGINAL ANSWER

Maybe you could try to create one class per state :

class TelnetState

  include Telnet::Codes
  attr_reader :plain, :reply

  def initialize( plain = '', reply ='' )
    @plain = plain
    @reply = reply
  end

  def accept_char( char )
    raise "Abstract Method - not implemented"
  end

end

class TelnetIDLEState < TelnetState

  def accept_char( char )
    if char == IAC 
      TelnetIACState.new( plain, reply )
    else
      @plain << char
      self        
    end
  end

end

class TelnetIACState < TelnetState

   def accept_char( char )
     update_plain!( char )
     next_state( char )
   end

   private

   def update_plain!( char )
     return if [WILL, WONT, DO, DONT, IP,DM].include? char
     @plain << char == IAC ? char : IAC + char
   end

   def next_state( char ) 
     next = case char
              when WILL then TelnetWILLState
              when WONT then TelnetWONTState
              when DO   then TelnetDOState
              when WONT then TelnetDONTState
            end
     next ? next.new( plain, reply ) : self
   end

 end

... and so on, then :

class Telnet

  def telnet_state_machine( command )
    state = command.each_char.inject( TelnetIDLEState.new ) do |state, char|
              state.accept_char( char )
            end
    @plain, @reply = state.plain, state.reply
  end

end

Of course, there is certainly a way to improve that... just an idea. (I'm not really familiar with telnet protocol, either - so I did this a bit blindfolded)

Answer 2

As suggested by @m_x, This solution uses a parser of sorts, driven by the StringScanner built-in class. This is very compact, pretty readable, and gets rid of the state machine altogether:

Some methods to handle actions:

def accept(scanner)   
  @plain << scanner[1]
end

def reply_dont(scanner)
  @reply << IAC + DONT + scanner[1]
end

def reply_wont(scanner)
  @reply << IAC + WONT + scanner[1]
end

def ignore(scanner)
end

A list of telnet sequences:

# Telnet sequences to handle, and how to handle them

SEQUENCES = [
  [/#{IAC}(#{IAC})/, :accept],
  [/#{IAC}#{WILL}(.)/m, :reply_dont],
  [/#{IAC}#{WONT}(.)/m, :ignore],
  [/#{IAC}#{DO}(.)/m, :reply_wont],
  [/#{IAC}#{DONT}(.)/m, :ignore],
  [/#{IAC}#{IP}/, :ignore],
  [/#{IAC}#{DM}/, :ignore],
  [/(.)/m, :accept],
]

And the parser that uses them:

# Parse the the command.  Sets @plain and @reply

def parse_command(command)
  @plain = ''
  @reply = ''
  scanner = StringScanner.new(command)
  while !scanner.eos?
    SEQUENCES.each do |regexp, method|
      if scanner.scan(regexp)
        send method, scanner
        break
      end
    end
  end
end

Answer 3

As suggested, I tried using a gem like state_machine. Well, not a gem like state_machine, but that particular gem. The state machine defintion is:

    state_machine :state, :initial => :idle do

      event :iac do
        transition :idle => :iac
        transition :iac => :idle
        transition [:will, :wont, :do, :dont] => :idle
      end

      event :will do
        transition :iac => :will
        transition :will => :idle
      end

      event :dm do
        transition :iac => :idle
      end

      event :ip do
        transition :iac => :idle
      end

      event :dont do
        transition :iac => :dont
      end

      event :other do
        transition all => :idle
      end

      event :do do
        transition :iac => :do
      end

      event :wont do
        transition :iac => :wont
      end

      before_transition :from => :do, :to => :idle, :do => :send_wont
      before_transition :from => :will, :to => :idle, :do => :send_dont
      before_transition :from => :idle, :to => :idle, :do => :accept_plain
      before_transition :from => :iac, :to => :idle, :on => :other, :do => :accept_unknown_iac
      before_transition :from => :iac, :to => :idle, :on => :iac, :do => :accept_plain

    end

The state_machine DSL is clean and relatively compact. Unfortunately, it is organized around events that cause transitions, rather than states that have transitions. That's not bad, but it's completely inside-out to how I think about state machines. It also makes you define actions separately from the transitions. The net result is that I can't see the state machine from this definition.

However, the state_machine gem has a rake task to generate a state diagram directly from the code. Here's what it generated for the above code:

Other than the lack of actions, the state diagram generated by state_machine is very nice.

Answer 4

Here's a refactoring that extracts the definition of the state machine into a compact data structure. First a few lambdas for the actions:

    # State machine actions.
    #
    # @param plain [String] Accumulated plaintext
    # @param reply [String] Accumulated telnet reply
    # @param c [String] The current character

    EMIT_PLAIN = lambda do |plain, reply, c|
      plain << c
    end

    EMIT_PLAIN_UNKNOWN_SEQUENCE = lambda do |plain, reply, c|
      plain << IAC + c
    end

    EMIT_REPLY_DONT = lambda do |plain, reply, c|
      reply << IAC + DONT + c
    end

    EMIT_REPLY_WONT = lambda do |plain, reply, c|
      reply << IAC + WONT + c
    end

Then the state machine definition itself. Of all of the state-machine approaches I tried, this one has the most succinct definition:

    # The definition of the state machine used to recognize and handle
    # Telnet sequences.
    #
    # This is organized as nested hashes and arrays:
    #
    # The outer, or "state" hash is:
    # * key: state
    # * value: character hash
    #
    # The character hash is:
    # * key: character code, or :else
    # * value: action array
    #
    # The action array is a list of zero or more:
    # * lambda - The action to perform
    # * symbol - the next state

    STATES = {
      :idle => {
        IAC => [:iac],
        :else => [EMIT_PLAIN],
      },
      :iac => {
        IAC => [EMIT_PLAIN, :idle],
        WILL => [:will],
        WONT => [:wont],
        DO => [:do],
        DONT => [:dont],
        IP => [:idle],
        DM => [:idle],
        :else => [EMIT_PLAIN_UNKNOWN_SEQUENCE, :idle],
      },
      :will => {
        :else => [EMIT_REPLY_DONT, :idle],
      },
      :wont => {
        :else => [:idle],
      },
      :do => {
        :else => [EMIT_REPLY_WONT, :idle],
      },
      :dont => {
        :else => [:idle]
      }
    }

And, finally, the functions that drive the state machine:

    # Parse the the command.  Sets @plain and @reply
    #
    # @param command [String] The command to parse

    def parse_command(command)
      @plain = ''
      @reply = ''
      @state = :idle
      command.each_char do |c|
        character_hash = STATES[@state]
        raise "Unknown state #{@state.inspect}" unless character_hash
        actions = character_hash[c] || character_hash[:else]
        raise "Missing :else" unless actions
        apply_actions actions, c
      end
    end

    def apply_actions(actions, c)
      actions.each do |action|
        apply_action action, c
      end
    end

    def apply_action(action, c)
      if action.is_a?(Symbol)
        @state = action
      elsif action.respond_to?(:call)
        action.call @plain, @reply, c
      else
        raise "Unknown action #{action.inspect}"
      end
    end

This does a good job of showing the state machine in a compact and clear format, but the machinery that drives the state machine (parse_command and friends) seem like clutter. The state machine definition may be too compact--the amount of comments I felt necessary to describe it are a sign of that.