This is how I've done it.

func application(_ application: UIApplication, didFinishLaunchingWithOptions launchOptions: [UIApplication.LaunchOptionsKey: Any]?) -> Bool {
    application.registerForRemoteNotifications()
    Messaging.messaging().delegate = self
    return true
}

func application(_ application: UIApplication, didRegisterForRemoteNotificationsWithDeviceToken deviceToken: Data) {
    Messaging.messaging().apnsToken = deviceToken;
}

func messaging(_ messaging: Messaging, didReceiveRegistrationToken fcmToken: String?) {
    UserDefaults.standard.set(fcmToken!, forKey: "fcmToken")
}