Two-Factor Authentication With Vapor

Retrieving OTP Tokens

Next up, add a function to create a TwoFactorToken.Public instance to return to users:

func asPublic() -> Public {
  let issuer = "DiningIn"
  let url = "otpauth://totp/\(self.user.username)?secret=\(self.key)&issuer=\(issuer)"
    return Public(
      backupCodes: self.backupTokens,
      key: self.key,
      label: self.user.username,
      issuer: issuer,
      url: url)
}

As the last step, add a static function to find a TwoFactorToken for a given User:

static func find(
  for user: User, 
  on db: Database
) -> EventLoopFuture<TwoFactorToken> {
  user.$twoFactorToken
    .query(on: db)
    .with(\.$user)
    .first()
    .unwrap(or: Abort(.internalServerError, reason: "No 2FA token found"))
}

This uses Fluent’s @Children property wrapper to query a user’s linked tokens, eager loads the user, and then unwraps and returns it. However, it doesn’t compile right yet, because a user doesn’t have a twoFactorToken property.

In User.swift, add the following to User under updatedAt:

@Children(for: \TwoFactorToken.$user)
var twoFactorToken: [TwoFactorToken]

Here, you complete the relationship between TwoFactorToken and User.

Validating OTP Tokens

While still in User.swift, add another @Field property under the passwordHash field to store whether or not a user has 2FA enabled:

@Field(key: "two_factor_enabled")
var twoFactorEnabled: Bool

Then, add the following line to the end of the User initializer to default twoFactorEnabled to false:

self.twoFactorEnabled = false

Adding the Migrations

Since these are database schema-related changes, the next step is to create.

Add a new file in the Migrations folder called CreateTwoFactorToken.swift with the following contents:

import Fluent

struct CreateTwoFactorToken: Migration {
  func prepare(on database: Database) -> EventLoopFuture<Void> {
    // 1
    let create = database.schema(TwoFactorToken.schema)
      .id()
      .field("user_id", .uuid, .required, .references(User.schema, "id"))
      .field("key", .string, .required)
      .unique(on: "user_id", "key")
      .field("backup_tokens", .array(of: .string), .required)
      .create()

    // 2    
    let update = database.schema(User.schema)
      .field("two_factor_enabled", .bool, .sql(.default(false)))
      .update()
    
    // 3
    return create.and(update).transform(to: ())
  }
  
  func revert(on database: Database) -> EventLoopFuture<Void> {
    // 4
    let delete = database.schema(TwoFactorToken.schema)
      .delete()
    
    // 5
    let update = database.schema(User.schema)
      .deleteField("two_factor_enabled")
      .update()

    // 6
    return delete.and(update).transform(to: ())
  }
}

This migration will do the following:

1. To prepare, first create the schema that creates two_factor_tokens.
2. Then, add two_factor_enabled to users.
3. Using EventLoopFuture.and(_:), combine the two futures into a single future and transform it to Void.
4. To revert, first create a delete step for two_factor_tokens.
5. Then, create a step to remove two_factor_field from users.
6. Again, using EventLoopFuture.and(_:), combine the two steps and transform to Void.

In configure.swift, add the migration. Below all other migrations, add the following:

app.migrations.add(CreateTwoFactorToken())

This ensures running the app will create the table and fields required for 2FA.

Setting up the Handshake

Enabling 2FA involves a handshake between client and server. This process consists of the following steps:

1. The user will request for 2FA to be enabled.
2. The API will generate a token and send it back in the response.
3. The user will add the token to their 2FA app, like Authy or Google Authenticator.
4. The user will send a code generated by Authy or Google Authenticator to the API.
5. The API will verify the submitted code and enable 2FA for the user, if valid.

To facilitate this 2FA handshake, you’ll add two endpoints in UserController.swift: one to generate a 2FA token, and one to validate the 2FA flow.

Below boot, add the following:

// 1
fileprivate func getTwoFactorToken(
  _ req: Request
) throws -> EventLoopFuture<TwoFactorToken.Public> {
  // 2
  let user = try req.auth.require(User.self)
  // 3
  func _token() -> EventLoopFuture<TwoFactorToken.Public> {
    TwoFactorToken.find(for: user, on: req.db).map { $0.asPublic() }
  }
  // 4
  if user.twoFactorEnabled {
    // 5
    return _token()
  } else {
    // 6
    return user.$twoFactorToken.get(on: req.db)
      .flatMapThrowing { token -> TwoFactorToken? in
        // 7
        if let _ = token.first {
          return nil
        }
        // 8
        let token = try TwoFactorToken.generate(for: user)
        return token
      }
      // 9
      .optionalFlatMap { $0.save(on: req.db) }
      // 10
      .flatMap { _ in _token() }
  }
}

Here’s what’s going on here:

1. Create a request handler returning TwoFactorToken.Public, which was created earlier.
2. Get an instance of User by requiring authentication.
3. Define a function that will find a TwoFactorToken for a user and return its public type.
4. Check if the user already has 2FA enabled.
5. If 2FA is already enabled, return the existing 2FA token’s information.
6. Otherwise, if 2FA has not been set up, try to find an existing but unverified TwoFactorToken.
7. If you find a token, return nil. This might seem counter intuitive, but it allows you to skip step 9.
8. When you don’t find a token, generate a token and return it.
9. Using optionalFlatMap(_:), save the generated token. This code will only execute when step 6 returns a non-nil value.
10. Finally, use _token() to return the public version of a user’s token.

This takes care of the first three steps of the 2FA handshake.

Below getTwoFactorToken, add the following function that will take care of the last part of the handshake:

// 1
fileprivate func enableTwoFactor(
  _ req: Request
) throws -> EventLoopFuture<HTTPStatus> {
  // 2
  let user = try req.auth.require(User.self)
  // 3
  if user.twoFactorEnabled {
    return req.eventLoop.makeSucceededFuture(.ok)
  }
  // 4
  guard let t = req.headers.first(name: "X-Auth-2FA") else {
    throw Abort(.badRequest)
  }
  // 5
  return TwoFactorToken.find(for: user, on: req.db).flatMap { token in
    // 6
    guard token.validate(t, allowBackupCode: false) else {
      // 7
      return token.delete(force: true, on: req.db).flatMapThrowing {
        throw Abort(.unauthorized)
      }
    }
    // 8
    user.twoFactorEnabled = true
    return user.save(on: req.db).transform(to: .ok)
  }
}

Here’s what’s happening in the code above:

1. Create a request handler returning a HTTP status code.
2. Get an instance of the User type by requiring authentication.
3. Check if the user already has 2FA enabled. If so, immediately return 200 OK.
4. When 2FA isn’t enabled, attempt to get the submitted code from the X-Auth-2FA header. Return 400 Bad Request if the header is missing.
5. Get the user’s 2FA token from the database.
6. Make sure the code is valid. Disallow backup codes here to be 100% sure the submitted code checks that the HMAC key has been saved correctly.
7. If the code isn’t valid, return 401 Unauthorized and delete the 2FA token from the database.
8. If the code is valid, enable 2FA for the user, save it to the database and return 200 OK

In boot, add the following two lines below the tokenProtected line to register these two routes to the router:

tokenProtected.get("me", "twofactortoken", use: getTwoFactorToken)
tokenProtected.post("me", "enabletwofactor", use: enableTwoFactor)

Now it’s time to test the handshake!