/*
 * uIAM - the IAM for microservices
 * Copyright (C) SoftInstigate Srl
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 * 
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package io.uiam.handlers.injectors;

import io.uiam.handlers.PipedHttpHandler;
import io.uiam.handlers.exchange.ByteArrayRequest;
import io.uiam.plugins.authentication.impl.BaseAccount;
import io.undertow.security.idm.Account;
import io.undertow.server.HttpServerExchange;
import io.undertow.util.HttpString;
import java.util.List;
import java.util.Map;

/**
 * Adds the following X-Forwarded custom headers to the proxied request:
 *
 * 'X-Forwarded-Account-Id', 'X-Forwarded-Account-Roles' and other headers set
 * with Response.addXForwardedHeader()
 *
 * @author Andrea Di Cesare {@literal <andrea@softinstigate.com>}
 *
 */
public class XForwardedHeadersInjector extends PipedHttpHandler {

    /**
     * Creates a new instance of AccountHeadersInjector
     *
     * @param next
     */
    public XForwardedHeadersInjector(PipedHttpHandler next) {
        super(next);
    }

    /**
     * Creates a new instance of AccountHeadersInjector
     *
     */
    public XForwardedHeadersInjector() {
        super(null);
    }

    /**
     *
     * @param exchange
     * @throws Exception
     */
    @Override
    public void handleRequest(HttpServerExchange exchange) throws Exception {
        if (exchange != null && exchange.getSecurityContext() != null
                && exchange.getSecurityContext()
                        .getAuthenticatedAccount() != null) {
            Account a = exchange.getSecurityContext().getAuthenticatedAccount();

            if (a.getPrincipal() != null) {
                exchange.getRequestHeaders()
                        .add(getHeaderForAccountId(),
                                a.getPrincipal().getName());
            }

            if (a instanceof BaseAccount && ((BaseAccount) a).getRoles() != null) {
                ((BaseAccount) a).getRoles()
                        .stream().forEachOrdered(role -> exchange
                        .getRequestHeaders()
                        .add(getHeaderForAccountRoles(), role));

            }

            Map<String, List<String>> xfhs = ByteArrayRequest.wrap(exchange)
                    .getXForwardedHeaders();

            if (xfhs != null) {
                xfhs.entrySet().stream()
                        .forEachOrdered(m -> {
                            m.getValue().stream().forEachOrdered(v
                                    -> exchange.getRequestHeaders()
                                            .add(getHeaderForXH(m.getKey()),
                                                    v));
                        });
            }
        }

        next(exchange);
    }

    private HttpString getHeaderForXH(String suffix) {
        return HttpString.tryFromString("X-Forwarded-".concat(suffix));
    }

    private HttpString getHeaderForAccountId() {
        return getHeaderForXH("Account-Id");
    }

    private HttpString getHeaderForAccountRoles() {
        return getHeaderForXH("Account-Roles");
    }
}

@@ -2323,34 +2323,45 @@
 .add(get
-HeaderForAccountId
+XForwardedAccountIdHeaderName
 (),%0A    
@@ -2704,37 +2704,41 @@
 .add(get
-HeaderForAccountRoles
+XForwardedRolesHeaderName
 (), role
@@ -3187,27 +3187,36 @@
 .add(get
+XForwarded
 Header
-ForXH
+Name
 (m.getKe
@@ -3359,38 +3359,44 @@
 e);%0A    %7D%0A%0A    p
-rivate
+ublic static
  HttpString getH
@@ -3394,27 +3394,36 @@
 ring get
+XForwarded
 Header
-ForXH
+Name
 (String 
@@ -3508,38 +3508,44 @@
 ));%0A    %7D%0A%0A    p
-rivate
+ublic static
  HttpString getH
@@ -3547,26 +3547,37 @@
  get
-HeaderForAccountId
+XForwardedAccountIdHeaderName
 () %7B
@@ -3591,35 +3591,44 @@
   return get
+XForwarded
 Header
-ForXH
+Name
 (%22Account-Id
@@ -3647,14 +3647,20 @@
    p
-rivate
+ublic static
  Htt
@@ -3674,29 +3674,33 @@
  get
-HeaderForAccountRoles
+XForwardedRolesHeaderName
 () %7B
@@ -3722,19 +3722,28 @@
  get
+XForwarded
 Header
-ForXH
+Name
 (%22Ac

set public access to methods to get X-Forwarded header names