golangでunix domain socket経由で通信するechoサーバーを書いてみた

ssh-agent的に鍵を保持してくれるプログラムが欲しくて練習のためまずはecho serverを書いてみました。
golangを選んだ理由は以下です。

• bashからではunix domain socketを直接扱えないようなので諦めた
• OSX, Linuxで動いて欲しい
• クロスコンパイルが簡単(らしい)

昨日初めてgolangを書いたので、ツッコミ大歓迎というかむしろ色々指摘してもらえると嬉しいです。

unix domain socketの作成

ssh-agentみたいな感じということで、socketファイルは外部から指定というより起動時に自動で設定します。
mktemp -dのような動きをしてくれる関数を探したところ、ioutil.TempDirを見つけたので利用しています。
また、実行者以外からアクセスされないようにアクセス権を変更しています。

main.go

main() {
  tempDir, err := ioutil.TempDir("", "golang-sample-echo-server.")
  pid := strconv.Itoa(os.Getpid())
  socket := tempDir + "/server." + pid
  listener, err := net.Listen("unix", socket)
  if err != nil {
    log.Printf("error: %v\n", err)
    return
  }
  if err := os.Chmod(socket, 0700); err != nil {
    log.Printf("error: %v\n", err)
    return
  }
}

処理内容

1. ioutil.TempDir()でテンポラリディレクトリ作成
2. pidを取得して作成したテンポラリディレクトリ内にsocketファイルを作成
3. socketファイルの権限を所有者のみのアクセスに限定

待ち受け

上で作ったlistenerを受け取る関数を定義してmain関数から呼び出します。
コードはUnix Sockets in Go (stack oveflow)を参考にさせてもらっています。

main.go

func server(listener net.Listener) {
  for {
    fd, err := listener.Accept()
    if err != nil {
      return
    }
    go process(fd)
  }
}
func process(fd net.Conn) {
  for {
    buf := make([]byte, 512)
    nr, err := fd.Read(buf)
    if err != nil {
      break
    }
    data := buf[0:nr]
    fmt.Printf("Recieved: %v", string(data));
    _, err = fd.Write(data)
    if err != nil {
      log.Printf("error: %v\n", err)
      break
    }
  }
  fd.Close()
}

処理内容

1. listener.Accept()で入力待ち
2. 入力があるとprocess()を呼び出し一定サイズのバッファに読み込みながらechoする
   読み込むデータがなくなったらfdを閉じて終了
3. goroutineでprocess()を呼び出しているため、処理の戻りを待たずに再度listener.Accept()で入力待ち
4. listenerがcloseされた場合はlistener.Accept()がエラーを返すので、forループを抜ける

終了時の掃除

ここまでの処理では終了時にゴミが残ってしまうので、適宜掃除を行います。

コードはGolang で書いた Web アプリケーションを UNIX ドメインソケットで公開を参考にさせてもらいました。

main.go

func shutdown(listener net.Listener, tempDir string, close chan int) {
  c := make(chan os.Signal, 2)
  signal.Notify(c, os.Interrupt, syscall.SIGTERM)
  go func() {
    interrupt := 0
    for {
      s := <-c
      switch s {
      case os.Interrupt:
        if (interrupt == 0) {
          fmt.Println("Interrupt...")
          interrupt++
          continue
        }
      }
      break
    }
    if err := listener.Close(); err != nil {
      log.Printf("error: %v\n", err)
    }
    if err := os.Remove(tempDir); err != nil {
      log.Printf("error: %v\n", err)
    }
    close <- 1
  }()
}

処理内容

1. シグナルを待ち受け用にchannelを作成
2. goroutineでシグナル発生時の処理関数を呼び出し、channelの通知待ち状態にしておく
3. シグナル発生時にgoroutineで登録してある処理が再開
   • Ctrl+C の場合は初回は見逃すように設定
4. listener.Close(), ディレクトリの削除を行なう
5. 終了を呼び出し元から渡されたchannelに通知

最終的なmain関数

ビルド可能なコードはgithubにpushしてあります。

main.go

func main() {
  log.SetFlags(log.Lshortfile)
  tempDir, err := ioutil.TempDir("", "golang-sample-echo-server.")
  pid := strconv.Itoa(os.Getpid())
  socket := tempDir + "/server." + pid
  listener, err := net.Listen("unix", socket)
  if err != nil {
    log.Printf("error: %v\n", err)
    return
  }
  if err := os.Chmod(socket, 0700); err != nil {
    log.Printf("error: %v\n", err)
    return
  }
  close := make(chan int)
  shutdown(listener, tempDir, close)
  fmt.Printf("GOLANG_SAMPLE_SOCK=%v;export GOLANG_SAMPLE_SOCK;\n", socket)
  fmt.Printf("GOLANG_SAMPLE_PID=%v;export GOLANG_SAMPLE_PID;\n", pid)
  server(listener)
  _ = <-close
}

処理内容

unix domain socket作成からの続きです。

1. shutdown()関数に渡すchannelを作成、shutdown関数を呼び出し
2. クライアント側の接続先のためにsocketとpidを出力
3. server()関数に作成したlistenerを渡して起動
4. shutdown()の終了を待機
   1. Ctrl+Cを2回実行すると、shutdown()内のlistener.Close()が呼ばれる
   2. listener.Close()によりserver()内のlistener.Accept()がエラーを返却する
   3. listener.Accept()のエラーを受けてserver()のループがbreak
   4. shutdown側のchannelを待機

動作確認

$ golang/bin/golang-sample-echo-server
GOLANG_SAMPLE_SOCK=/var/folders/ld/92yggdj91l72jq205zr0l1540000gn/T/golang-sample-echo-server.660221887/server.61435;export GOLANG_SAMPLE_SOCK;
GOLANG_SAMPLE_PID=61435;export GOLANG_SAMPLE_PID;
Recieved: sample echo server test

$ echo 'sample echo server test' | nc -U /var/folders/ld/92yggdj91l72jq205zr0l1540000gn/T/golang-sample-echo-server.660221887/server.61435
sample echo server test

残タスク

• コマンドラインオプションの対応を追加
• 自らdaemon化するように処理を追加
• 自らを停止するオプション・処理の追加

おまけ

ログをわかりやすく

エラー発生時にログがどの行から出ているのがわかりづらくてとても困りました。
以下の行をmain関数の最初の行に追加することで[ファイル名]:[行番号]: [メッセージ]というフォーマットに変更してくれて助かりました。

log.SetFlags(log.Lshortfile)

2016/10/12追記

同僚にいくつか指摘されたので対応しました。

fd.Close()はdeferに変更

末尾にあったfd.Close()を先頭に持ってきてdeferをつけただけです。

main.go

func process(fd net.Conn) {
  defer fd.Close()
  for {
    buf := make([]byte, 512)
    nr, err := fd.Read(buf)
    if err != nil {
      break
    }
    data := buf[0:nr]
    fmt.Printf("Recieved: %v", string(data));
    _, err = fd.Write(data)
    if err != nil {
      log.Printf("error: %v\n", err)
      break
    }
  }
}

serverを構造体にする

• サーバーがListenerの管理をする為に、main関数内でやっていたことの一部をサーバー側に移動しています。

server.go

package main

import (
  "net"
  "log"
  "os"
  "fmt"
)

type Server struct {
  listener net.Listener
}

func NewServer() *Server{
  s := new(Server)
  return s;
}

func (s *Server) Open(socket string) {
  listener, err := net.Listen("unix", socket)
  if err != nil {
    log.Printf("error: %v\n", err)
    return
  }
  s.listener = listener;
  if err := os.Chmod(socket, 0700); err != nil {
    log.Printf("error: %v\n", err)
    s.Close()
    return
  }
}

func (s *Server) Close() {
  if err := s.listener.Close(); err != nil {
    log.Printf("error: %v\n", err)
  }
}

func (s *Server) Start() {
  for {
    fd, err := s.listener.Accept()
    if err != nil {
      return
    }
    go s.Process(fd)
  }
}

func (s *Server) Process(fd net.Conn) {
  defer fd.Close()
  for {
    buf := make([]byte, 512)
    nr, err := fd.Read(buf)
    if err != nil {
      break
    }
    data := buf[0:nr]
    fmt.Printf("Recieved: %v", string(data));
    _, err = fd.Write(data)
    if err != nil {
      log.Printf("error: %v\n", err)
      break
    }
  }
}

main.go

package main

import (
  "os"
  "os/signal"
  "io/ioutil"
  "fmt"
  "syscall"
  "strconv"
  "log"
)

func main() {
  log.SetFlags(log.Lshortfile)
  tempDir, err := ioutil.TempDir("", "golang-sample-echo-server.")
  if err != nil {
    log.Printf("error: %v\n", err)
    return
  }
  pid := strconv.Itoa(os.Getpid())
  socket := tempDir + "/server." + pid
  if err := os.Chmod(tempDir, 0700); err != nil {
    log.Printf("error: %v\n", err)
    return
  }
  defer func() {
    if err := os.Remove(tempDir); err != nil {
      log.Printf("error: %v\n", err)
    }
  }()
  server := NewServer()
  server.Open(socket)
  registerShutdown(server)
  fmt.Printf("GOLANG_SAMPLE_SOCK=%v;export GOLANG_SAMPLE_SOCK;\n", socket)
  fmt.Printf("GOLANG_SAMPLE_PID=%v;export GOLANG_SAMPLE_PID;\n", pid)
  server.Start()
}

func registerShutdown(server *Server) {
  c := make(chan os.Signal, 2)
  signal.Notify(c, os.Interrupt, syscall.SIGTERM)
  go func() {
    interrupt := 0
    for {
      s := <-c
      switch s {
      case os.Interrupt:
        if (interrupt == 0) {
          fmt.Println("Interrupt...")
          interrupt++
          continue
        }
      }
      break
    }
    server.Close()
  }()
}

2016/10/13追記

(動かなかったので2016/10/14に微調整しました)
更にエラー処理について指摘されたので対応しました。
差分だけの方がわかりやすそうなので差分のみです。

Serverはエラー発生時にErrorを返却し、実際のエラーの処理はmain内で行うようにしました。

diff --git a/main.go b/main.go
index 6478944..92dbdbf 100644
--- a/main.go
+++ b/main.go
@@ -29,11 +29,14 @@ func main() {
     }
   }()
   server := NewServer()
-  server.Open(socket)
+  if err := server.Open(socket); err != nil {
+    log.Printf("error: %v\n", err)
+    return;
+  }
   registerShutdown(server)
   fmt.Printf("GOLANG_SAMPLE_SOCK=%v;export GOLANG_SAMPLE_SOCK;\n", socket)
   fmt.Printf("GOLANG_SAMPLE_PID=%v;export GOLANG_SAMPLE_PID;\n", pid)
-  server.Start()
+  server.Start();
 }

 func registerShutdown(server *Server) {
@@ -53,6 +56,8 @@ func registerShutdown(server *Server) {
       }
       break
     }
-    server.Close()
+    if err := server.Close(); err != nil {
+      log.Printf("error: %v\n", err)
+    }
   }()
 }
diff --git a/server.go b/server.go
index 64a03af..3f40115 100644
--- a/server.go
+++ b/server.go
@@ -2,7 +2,6 @@ package main

 import (
   "net"
-  "log"
   "os"
   "fmt"
 )
@@ -16,37 +15,37 @@ func NewServer() *Server{
   return s;
 }

-func (s *Server) Open(socket string) {
+func (s *Server) Open(socket string) error {
   listener, err := net.Listen("unix", socket)
   if err != nil {
-    log.Printf("error: %v\n", err)
-    return
+    return err
   }
   s.listener = listener;
   if err := os.Chmod(socket, 0600); err != nil {
-    log.Printf("error: %v\n", err)
     s.Close()
-    return
+    return err
   }
+  return nil
 }

-func (s *Server) Close() {
+func (s *Server) Close() error{
   if err := s.listener.Close(); err != nil {
-    log.Printf("error: %v\n", err)
+    return err;
   }
+  return nil
 }

 func (s *Server) Start() {
   for {
     fd, err := s.listener.Accept()
     if err != nil {
-      return
+      break;
     }
     go s.Process(fd)
   }
 }

-func (s *Server) Process(fd net.Conn) {
+func (s *Server) Process(fd net.Conn) error{
   defer fd.Close()
   for {
     buf := make([]byte, 512)
@@ -58,8 +57,8 @@ func (s *Server) Process(fd net.Conn) {
     fmt.Printf("Recieved: %v", string(data));
     _, err = fd.Write(data)
     if err != nil {
-      log.Printf("error: %v\n", err)
-      break
+      return err
     }
   }
+  return nil

参考させてもらったサイト

The Go Programming Language
- Packege net
Golang で書いた Web アプリケーションを UNIX ドメインソケットで公開
Unix Sockets in Go(stack oveflow)