#78 GoでTOTP(Time-based One-Time Password)を実装してみる

2023 / 2 / 1

2要素認証で使われるTOTP(Time-based One-Time Password)をGoで実装してみました。

今回はこちらでやったことを解説してみます。

実装してみたもの

以下の手順っぽいことができるサーバーを書いてみました。本来はパスワードの保存が必要ですが、本質ではないので省略してます。

ユーザー登録
ログイン
Google認証システムで登録可能なQRコード作成
コードを使って2要素認証の設定完了

ユーザー登録

まずユーザーのデータ定義をこんな感じにしてみます。

type User struct {
	Username        string
	TotpSecret      string
	SecretVerified  bool
	SetupTotpSecret string
}

Usernameはログインに使うユーザー名です。メールアドレスでログインするサービスだったらEmailになるでしょう。TotpSecretは認証に使うシークレットです。今回の範囲では使いません。

/api/signupを叩くとユーザーテーブルにユーザーを追加するようにしてみます。

m.HandleFunc("/api/signup", signup(models)).Methods("POST")

func signup(models *model.Models) http.HandlerFunc {
  return func(w http.ResponseWriter, r *http.Request) {
    bodyBytes, err := io.ReadAll(r.Body)
    if err != nil {
      return
    }
    var data map[string]interface{}
    err = json.Unmarshal(bodyBytes, &data)
    if err != nil {
      writeError(w, http.StatusBadRequest, "request body is not JSON")
      return
    }
    username, _ := data["username"].(string)
    if len(username) == 0 {
      writeError(w, http.StatusBadRequest, "username must not be empty")
      return
    }
    // データベースに追加する
    // ちゃんとしたサービスならパスワードもbcryptあたりで保存する
    u, err := models.User.Create(username)
    if err != nil {
      writeError(w, http.StatusInternalServerError, err.Error())
      return
    }
    writeMap(w, u.ToMap())
  }
}

サンプル実装はメモリ上に保存するようにしています。

m.HandleFunc("/api/login", login(models)).Methods("POST")

func login(models *model.Models) http.HandlerFunc {
  return func(w http.ResponseWriter, r *http.Request) {
    bodyBytes, err := io.ReadAll(r.Body)
    if err != nil {
      return
    }
    var data map[string]interface{}
    err = json.Unmarshal(bodyBytes, &data)
    if err != nil {
      writeError(w, http.StatusBadRequest, "request body is not JSON")
      return
    }
    username, _ := data["username"].(string)
    if len(username) == 0 {
      writeError(w, http.StatusBadRequest, "username must not be empty")
      return
    }
    // 本来はパスワードも渡す
    token, err := models.User.Login(username)
    if err != nil {
      writeError(w, http.StatusInternalServerError, err.Error())
      return
    }
    writeMap(w, map[string]interface{}{
      "token": token,
    })
  }
}

QRコード作成

ここからが本題です。今回はUIが無いのでアレですが、ログイン後、2要素認証がまだ有効になっていないユーザーだったらQRコードを表示し、Google認証システムアプリで読み込んでもらいます。

React等のSPAで使われる想定で、APIを叩くとQRコードが返ってくるようにします。

m.HandleFunc("/api/totp/setup", totpSetup(models)).Methods("POST")

func totpSetup(models *model.Models) http.HandlerFunc {
  return func(w http.ResponseWriter, r *http.Request) {
    authorization := r.Header.Get("Authorization")
    authValues := strings.SplitN(authorization, " ", 2)
    jwtToken := authValues[1]
  
    session, err := models.User.GetSession(jwtToken)
    if err != nil {
      fmt.Printf("getsession error %s\n", err)
      writeError(w, http.StatusBadRequest, "wrong token")
      return
    }
    // セッションからユーザー名を取得
    // Google認証アプリに表示される名前になる点注意
    username := session.Username
    // QRコードの画像を作る
    img, err := models.User.CreateTotpCode(username)
    if err != nil {
      writeError(w, http.StatusInternalServerError, err.Error())
      return
    }
    err = png.Encode(w, img)
    if err != nil {
      fmt.Printf("Failed to encode %s", err)
      return
    }
  }
}

QRコードを作る部分はこんな感じです。

func (r *repo) CreateTotpCode(username string) (image.Image, error) {
  // ストレージからユーザー情報をとってくる
	u, err := r.GetByUsername(username)
	if err != nil {
		return nil, err
	}
	// keyを作る
	key, err := totp.Generate(totp.GenerateOpts{
		Issuer:      "example.mokelab.com",
		AccountName: username,
		Period:      30,
		SecretSize:  20,
		Secret:      []byte{},
		Digits:      otp.DigitsSix,
		Algorithm:   otp.AlgorithmSHA1,
		Rand:        rand.Reader,
	})
	if err != nil {
		return nil, err
	}
	// QRコードの画像を作る
	img, err := key.Image(200, 200)
	if err != nil {
		return nil, err
	}
	// 有効にする作業中のシークレットとして保存しておく
	u.SetupTotpSecret = key.Secret()
	r.users[username] = u
	return img, nil
}

keyの生成はgithub.com/pquerna/otpの関数をほぼサンプルの通りに呼ぶだけです。Issuerは発行者名でAccountNameはユーザー名を指定します。この2つはGoogle認証システムのアプリ上で表示されます。

keyを作ったら、後の確認フェーズでシークレットを使うのでセットアップ時のシークレットとしてストレージに保存しておきます。

最後にQRコードの画像を返します。これをユーザーに表示し、Google認証アプリで登録してもらいます。

コードの確認をする

QRコードを認証アプリに登録してもらったら、確認のためにコードをサーバーに送信してもらいます。

セッションからユーザー名をとったり、リクエストからコードをとったりする部分があるのでやや長いですが、大事なのはVerifySetupTotpCode()の部分だけです。

m.HandleFunc("/api/totp/setup/verify", totpSetupVerify(models)).Methods("POST")

func totpSetupVerify(models *model.Models) http.HandlerFunc {
  return func(w http.ResponseWriter, r *http.Request) {
    authorization := r.Header.Get("Authorization")
    authValues := strings.SplitN(authorization, " ", 2)
    jwtToken := authValues[1]
  
    session, err := models.User.GetSession(jwtToken)
    if err != nil {
      fmt.Printf("getsession error %s\n", err)
      writeError(w, http.StatusBadRequest, "wrong token")
      return
    }
  
    bodyBytes, err := io.ReadAll(r.Body)
    if err != nil {
      return
    }
    var data map[string]interface{}
    err = json.Unmarshal(bodyBytes, &data)
    if err != nil {
      writeError(w, http.StatusBadRequest, "request body is not JSON")
      return
    }
  
    username := session.Username
    code, _ := data["code"].(string)
    valid, err := models.User.VerifySetupTotpCode(username, code)
    if err != nil {
      writeError(w, http.StatusBadRequest, err.Error())
      return
    }
    if valid {
      writeMap(w, map[string]interface{}{
        "result": "verified",
      })
    } else {
      writeMap(w, map[string]interface{}{
        "result": "failed",
      })
    }
  }
}

VerifySetupTotpCode()では、totp.Validate()を呼ぶだけです。この際keyを作る時に取得したシークレットが必要となるのでストレージから取得しましょう。


func (r *repo) VerifySetupTotpCode(username string, passCode string) (bool, error) {
  u, err := r.GetByUsername(username)
  if err != nil {
    return false, err
  }
  if len(u.SetupTotpSecret) == 0 {
    return false, errors.New("not prepared")
  }
  valid := totp.Validate(passCode, u.SetupTotpSecret)
  if valid {
    u.TotpSecret = u.SetupTotpSecret
    u.SetupTotpSecret = ""
    return true, nil
  } else {
    return false, nil
  }
}

まとめ

GoでTOTPを実装してみました。keyを作った時点でシークレットを保存しないと後の検証で詰むのがポイントでしょうか。。。