func main() { kingpin.Parse() msgChan, err := buildStream() if err != nil || msgChan == nil { kingpin.FatalUsage("Failed to create input stream: " + err.Error()) } // Disable regular logging, cuz bonjour logs an error that isn't, and it's // confusing. dontlog.SetOutput(NullWriter{}) if *host == "" { *host = getLocalIP() } if *port == 0 { // TODO: run a server that sends the same text on TCP conn *port = 45897 } serv, err := service.New(*host, *port, msgChan) if err != nil || serv == nil { log.Panic("Failed to create service", "err", err) } // Watch for signal to clean up before we exit signals := make(chan os.Signal, 1) signal.Notify(signals, os.Interrupt, os.Kill) <-signals log.Info("Shutting down") serv.Stop() log.Info("Stopped service") }
func (s *Service) Stop() { defer s.wg.Wait() log.Info("Stopping service") close(s.stop) // TCPListener might still be blocked, so open a conn to it ourselves // to free it up. if conn, err := net.DialTCP("tcp", nil, s.addr); err != nil { log.Error("Failed to close down TCP listener", "err", err) } else { conn.Close() log.Info("Shut down TCP listener") } }
func New(host string, port int, msgChan <-chan string) (*Service, error) { s := &Service{ host: host, port: port, messages: msgChan, stop: make(chan struct{}), } if addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", host, port)); err != nil { return nil, err } else { s.addr = addr } listener, err := net.ListenTCP("tcp", s.addr) if err != nil { return nil, err } log.Info("Listening for TCP connections", "address", s.addr) s.wg.Add(1) go s.serveTCP(listener) s.wg.Add(1) go s.start() return s, nil }
func CSVField(fieldIndex int, source DataSource) MessageBuilder { c := make(chan []string) go func() { for { reader := csv.NewReader(strings.NewReader(<-source)) reader.TrimLeadingSpace = true reader.FieldsPerRecord = -1 var values []string // Ignore first line - comment reader.Read() for { if record, err := reader.Read(); err == io.EOF { break } else if err != nil { log.Error("Error reading CSV data", "err", err) } else if len(record) <= fieldIndex { log.Error("Error reading CSV data: index out of bounds", "# fields in record", record, "index", fieldIndex) } else { values = append(values, record[fieldIndex]) } } if len(values) == 0 { log.Error("Didn't get any values from CSV") } else { log.Info("Read CSV data", "num values", len(values)) c <- values } } }() return c }
func (s *Service) stopBonjour(bonj *bonjour.Server) { if bonj == nil { return } s.wg.Add(1) go func() { defer s.wg.Done() log.Info("Shutting down bonjour service") bonj.Shutdown() // I guess bonjour wants us to wait some unspecied // amount? This is what blocking or channels are for :/ waitTime := time.Second * 5 log.Info("Waiting for bonjour service to clean itself up", "waitTime", waitTime) time.Sleep(waitTime) }() }
func (s *Service) start() { defer s.wg.Done() var bonj *bonjour.Server defer func(b **bonjour.Server) { s.stopBonjour(*b) }(&bonj) for { select { case msg := <-s.messages: if msg != s.currentMsg { s.currentMsg = msg s.stopBonjour(bonj) bonj = nil if msg != "" { log.Info("Registering service", "name", msg, "host", s.host, "port", s.port) var err error bonj, err = bonjour.RegisterProxy( msg, "_afpovertcp._tcp", "local", s.port, s.host, s.host, nil, nil) if err != nil || bonj == nil { log.Error("Failed to register service with bonjour", "err", err) bonj = nil } } } case <-s.stop: log.Info("Ending bonjour-updating routine") return } } }
func buildStream() (<-chan string, error) { var source inputs.DataSource var err error if *url != "" { log.Info("Reading messages from a url", "url", *url) source, err = inputs.Download(*url) if err != nil { return nil, err } } else if *file != "" { log.Info("Reading messages from file", "file", *file) source, err = inputs.FileWatcher(*file) if err != nil { return nil, err } } else if len(*say) > 0 { log.Info("Using a static message", "msg", *say) source, err = inputs.StaticText(strings.Join(*say, " ")) if err != nil { return nil, err } } if source == nil { return nil, errors.New("No source of data specified") } var builder inputs.MessageBuilder if *csvField >= 0 { log.Info("Iterating csv values", "field", *csvField) builder = inputs.CSVField(*csvField, source) } else if *words { log.Info("Iterating words") builder = inputs.WordGroups(source) } else { log.Info("Iterating lines") builder = inputs.Lines(source) } var chooser inputs.MessageChooser if *random { log.Info("Randomizing order") chooser = inputs.RandomMessageChooser(builder) } else { log.Info("Sequential order") chooser = inputs.SequentialMessageChooser(builder) } msgChan := inputs.RateLimit(*interval, inputs.LimitSize(40, inputs.Cleanup( chooser))) // Optional filters if *prefix != "" { msgChan = inputs.Prefix(*prefix, msgChan) } if *leet { msgChan = inputs.LeetSpeak(msgChan) } else if *mixedCase { msgChan = inputs.MixedCase(msgChan) } else if *lower { msgChan = inputs.LowerCase(msgChan) } else if *upper { msgChan = inputs.UpperCase(msgChan) } return msgChan, nil }