content: /*
   * Copyright 1999-2019 Alibaba Group Holding Ltd.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *      http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package com.alibaba.csp.sentinel.datasource.spring.cloud.config;
  
  import com.alibaba.csp.sentinel.log.RecordLog;
  
  import org.springframework.cloud.bootstrap.config.PropertySourceLocator;
  import org.springframework.cloud.config.client.ConfigClientProperties;
  import org.springframework.cloud.config.client.ConfigClientStateHolder;
  import org.springframework.cloud.config.environment.Environment;
  import org.springframework.cloud.config.environment.PropertySource;
  import org.springframework.core.annotation.Order;
  import org.springframework.core.env.CompositePropertySource;
  import org.springframework.core.env.MapPropertySource;
  import org.springframework.http.HttpEntity;
  import org.springframework.http.HttpHeaders;
  import org.springframework.http.HttpMethod;
  import org.springframework.http.HttpRequest;
  import org.springframework.http.HttpStatus;
  import org.springframework.http.MediaType;
  import org.springframework.http.ResponseEntity;
  import org.springframework.http.client.ClientHttpRequestExecution;
  import org.springframework.http.client.ClientHttpRequestInterceptor;
  import org.springframework.http.client.ClientHttpResponse;
  import org.springframework.http.client.SimpleClientHttpRequestFactory;
  import org.springframework.retry.annotation.Retryable;
  import org.springframework.util.Base64Utils;
  import org.springframework.util.StringUtils;
  import org.springframework.web.client.HttpClientErrorException;
  import org.springframework.web.client.HttpServerErrorException;
  import org.springframework.web.client.ResourceAccessException;
  import org.springframework.web.client.RestTemplate;
  
  import java.io.IOException;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  import static org.springframework.cloud.config.client.ConfigClientProperties.*;
  
  /**
   * <p>
   * {@link SentinelRuleLocator} which pulls Sentinel rules from remote server.
   * It retrieves configurations of spring-cloud-config client configurations from
   * {@link org.springframework.core.env.Environment}, such as {@code spring.cloud.config.uri=uri},
   * {@code spring.cloud.config.profile=profile}, and so on.
   * When rules are pulled successfully, it will be stored to {@link SentinelRuleStorage}.
   * </p>
   *
   * @author lianglin
   * @since 1.7.0
   */
  @Order(0)
  public class SentinelRuleLocator implements PropertySourceLocator {
  
      private RestTemplate restTemplate;
      private ConfigClientProperties defaultProperties;
      private org.springframework.core.env.Environment environment;
  
      public SentinelRuleLocator(ConfigClientProperties defaultProperties,
                                 org.springframework.core.env.Environment environment) {
          this.defaultProperties = defaultProperties;
          this.environment = environment;
      }
  
      /**
       * Responsible for pull data from remote server
       *
       * @param environment
       * @return correct data if success else a empty propertySource or null
       */
      @Override
      @Retryable(interceptor = "configServerRetryInterceptor")
      public org.springframework.core.env.PropertySource<?> locate(
          org.springframework.core.env.Environment environment) {
          ConfigClientProperties properties = this.defaultProperties.override(environment);
          CompositePropertySource composite = new CompositePropertySource("configService");
          RestTemplate restTemplate = this.restTemplate == null
              ? getSecureRestTemplate(properties)
              : this.restTemplate;
          Exception error = null;
          String errorBody = null;
          try {
              String[] labels = new String[] {""};
              if (StringUtils.hasText(properties.getLabel())) {
                  labels = StringUtils
                      .commaDelimitedListToStringArray(properties.getLabel());
              }
              String state = ConfigClientStateHolder.getState();
              // Try all the labels until one works
              for (String label : labels) {
                  Environment result = getRemoteEnvironment(restTemplate, properties,
                      label.trim(), state);
                  if (result != null) {
                      log(result);
                      // result.getPropertySources() can be null if using xml
                      if (result.getPropertySources() != null) {
                          for (PropertySource source : result.getPropertySources()) {
                              @SuppressWarnings("unchecked")
                              Map<String, Object> map = (Map<String, Object>)source
                                  .getSource();
                              composite.addPropertySource(
                                  new MapPropertySource(source.getName(), map));
                          }
                      }
                      SentinelRuleStorage.setRulesSource(composite);
                      return composite;
                  }
              }
          } catch (HttpServerErrorException e) {
              error = e;
              if (MediaType.APPLICATION_JSON.includes(e.getResponseHeaders().getContentType())) {
                  errorBody = e.getResponseBodyAsString();
              }
          } catch (Exception e) {
              error = e;
          }
          if (properties.isFailFast()) {
              throw new IllegalStateException(
                  "Could not locate PropertySource and the fail fast property is set, failing",
                  error);
          }
          RecordLog.warn("Could not locate PropertySource: " + (errorBody == null
              ? error == null ? "label not found" : error.getMessage()
              : errorBody));
          return null;
  
      }
  
      public org.springframework.core.env.PropertySource<?> refresh() {
          return locate(environment);
      }
  
      private void log(Environment result) {
  
          RecordLog.info("Located environment: name={}, profiles={}, label={}, version={}, state={}",
              result.getName(),
              result.getProfiles() == null ? "" : Arrays.asList(result.getProfiles()),
              result.getLabel(), result.getVersion(), result.getState());
  
          List<PropertySource> propertySourceList = result.getPropertySources();
          if (propertySourceList != null) {
              int propertyCount = 0;
              for (PropertySource propertySource : propertySourceList) {
                  propertyCount += propertySource.getSource().size();
              }
              RecordLog.info("[SentinelRuleLocator] Environment {} has {} property sources with {} properties",
                  result.getName(), result.getPropertySources().size(), propertyCount);
          }
      }
  
      private Environment getRemoteEnvironment(RestTemplate restTemplate,
                                               ConfigClientProperties properties, String label, String state) {
          String path = "/{name}/{profile}";
          String name = properties.getName();
          String profile = properties.getProfile();
          String token = properties.getToken();
          int noOfUrls = properties.getUri().length;
          if (noOfUrls > 1) {
              RecordLog.debug("[SentinelRuleLocator] Multiple Config Server Urls found listed.");
          }
  
          RecordLog.info("[SentinelRuleLocator] getRemoteEnvironment, properties={}, label={}, state={}",
              properties, label, state);
  
          Object[] args = new String[] {name, profile};
          if (StringUtils.hasText(label)) {
              if (label.contains("/")) {
                  label = label.replace("/", "(_)");
              }
              args = new String[] {name, profile, label};
              path = path + "/{label}";
          }
          ResponseEntity<Environment> response = null;
  
          for (int i = 0; i < noOfUrls; i++) {
              Credentials credentials = properties.getCredentials(i);
              String uri = credentials.getUri();
              String username = credentials.getUsername();
              String password = credentials.getPassword();
  
              RecordLog.info("[SentinelRuleLocator] Fetching config from server at: {}", uri);
  
              try {
                  HttpHeaders headers = new HttpHeaders();
                  addAuthorizationToken(properties, headers, username, password);
                  if (StringUtils.hasText(token)) {
                      headers.add(TOKEN_HEADER, token);
                  }
                  if (StringUtils.hasText(state) && properties.isSendState()) {
                      headers.add(STATE_HEADER, state);
                  }
  
                  final HttpEntity<Void> entity = new HttpEntity<>((Void)null, headers);
                  response = restTemplate.exchange(uri + path, HttpMethod.GET, entity,
                      Environment.class, args);
              } catch (HttpClientErrorException e) {
                  if (e.getStatusCode() != HttpStatus.NOT_FOUND) {
                      throw e;
                  }
              } catch (ResourceAccessException e) {
                  RecordLog.warn("[SentinelRuleLocator] ConnectTimeoutException on url <{}>."
                      + " Will be trying the next url if available", uri);
                  if (i == noOfUrls - 1) {
                      throw e;
                  } else {
                      continue;
                  }
              }
  
              if (response == null || response.getStatusCode() != HttpStatus.OK) {
                  return null;
              }
  
              Environment result = response.getBody();
              return result;
          }
  
          return null;
      }
  
      private RestTemplate getSecureRestTemplate(ConfigClientProperties client) {
          SimpleClientHttpRequestFactory requestFactory = new SimpleClientHttpRequestFactory();
          if (client.getRequestReadTimeout() < 0) {
              throw new IllegalStateException("Invalid Value for Read Timeout set.");
          }
          requestFactory.setReadTimeout(client.getRequestReadTimeout());
          RestTemplate template = new RestTemplate(requestFactory);
          Map<String, String> headers = new HashMap<>(client.getHeaders());
          if (headers.containsKey(AUTHORIZATION)) {
              // To avoid redundant addition of header
              headers.remove(AUTHORIZATION);
          }
          if (!headers.isEmpty()) {
              template.setInterceptors(Arrays.<ClientHttpRequestInterceptor>asList(
                  new GenericRequestHeaderInterceptor(headers)));
          }
  
          return template;
      }
  
      private void addAuthorizationToken(ConfigClientProperties configClientProperties,
                                         HttpHeaders httpHeaders, String username, String password) {
          String authorization = configClientProperties.getHeaders().get(AUTHORIZATION);
  
          if (password != null && authorization != null) {
              throw new IllegalStateException(
                  "You must set either 'password' or 'authorization'");
          }
  
          if (password != null) {
              byte[] token = Base64Utils.encode((username + ":" + password).getBytes());
              httpHeaders.add("Authorization", "Basic " + new String(token));
          } else if (authorization != null) {
              httpHeaders.add("Authorization", authorization);
          }
  
      }
  
      public void setRestTemplate(RestTemplate restTemplate) {
          this.restTemplate = restTemplate;
      }
  
      public static class GenericRequestHeaderInterceptor
          implements ClientHttpRequestInterceptor {
  
          private final Map<String, String> headers;
  
          public GenericRequestHeaderInterceptor(Map<String, String> headers) {
              this.headers = headers;
          }
  
          @Override
          public ClientHttpResponse intercept(HttpRequest request, byte[] body,
                                              ClientHttpRequestExecution execution) throws IOException {
              for (Map.Entry<String, String> header : headers.entrySet()) {
                  request.getHeaders().add(header.getKey(), header.getValue());
              }
              return execution.execute(request, body);
          }
  
          protected Map<String, String> getHeaders() {
              return headers;
          }
  
      }
  }
  